comm_tcp.h

Go to the documentation of this file.

/*****************************************************************************
* Multiscale Universal Interface Code Coupling Library                       *
*                                                                            *
* Copyright (C) 2019 Y. H. Tang, S. Kudo, X. Bian, Z. Li, G. E. Karniadakis  *
*                                                                            *
* This software is jointly licensed under the Apache License, Version 2.0    *
* and the GNU General Public License version 3, you may use it according     *
* to either.                                                                 *
*                                                                            *
* ** Apache License, version 2.0 **                                          *
*                                                                            *
* Licensed under the Apache License, Version 2.0 (the "License");            *
* you may not use this file except in compliance with the License.           *
* You may obtain a copy of the License at                                    *
*                                                                            *
* http://www.apache.org/licenses/LICENSE-2.0                                 *
*                                                                            *
* Unless required by applicable law or agreed to in writing, software        *
* distributed under the License is distributed on an "AS IS" BASIS,          *
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.   *
* See the License for the specific language governing permissions and        *
* limitations under the License.                                             *
*                                                                            *
* ** GNU General Public License, version 3 **                                *
*                                                                            *
* This program is free software: you can redistribute it and/or modify       *
* it under the terms of the GNU General Public License as published by       *
* the Free Software Foundation, either version 3 of the License, or          *
* (at your option) any later version.                                        *
*                                                                            *
* This program is distributed in the hope that it will be useful,            *
* but WITHOUT ANY WARRANTY; without even the implied warranty of             *
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the              *
* GNU General Public License for more details.                               *
*                                                                            *
* You should have received a copy of the GNU General Public License          *
* along with this program.  If not, see <http://www.gnu.org/licenses/>.      *
*****************************************************************************/
 
#ifndef COMM_TCP_H
#define COMM_TCP_H
 
 
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netdb.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <sys/epoll.h>
#include <sys/eventfd.h>
#include <fcntl.h>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <vector>
#include <deque>
#include <list>
#include <exception>
#include <atomic>
 
#include "message.h"
#include "comm.h"
#include "lib_uri.h"
#include "comm_factory.h"
 
namespace mui {
 
/*
 * smalluint is a integer which has variable length depend on
 * its value. smalluint requires small bytes if the value is smaller.
 * For example if the value is less than equal to 127, it only requires
 * one byte.
 * Format: the byte length is determined by the number of leading ones
 * of first byte.
 */
struct smalluint {
    uint64_t num;
    static int detect_size( unsigned char byte ){
        if(      byte <= 0x7fu ) return 1;
        else if( byte <= 0xbfu ) return 2;
        else if( byte <= 0xdfu ) return 3;
        else if( byte <= 0xefu ) return 4;
        else if( byte <= 0xf7u ) return 5;
        else if( byte <= 0xfbu ) return 6;
        else if( byte <= 0xfdu ) return 7;
        else if( byte <= 0xfeu ) return 8;
        else return 9;
    }
};
 
inline istream& operator>> ( istream& stream, smalluint& sml )
{
    unsigned char byte;
    stream >> byte;
    uint64_t num = 0u;
    int size;
    if(      byte <= 0x7fu ) { size = 1; num = byte&0xffu; }
    else if( byte <= 0xbfu ) { size = 2; num = byte&0x3fu; }
    else if( byte <= 0xdfu ) { size = 3; num = byte&0x1fu; }
    else if( byte <= 0xefu ) { size = 4; num = byte&0x0fu; }
    else if( byte <= 0xf7u ) { size = 5; num = byte&0x07u; }
    else if( byte <= 0xfbu ) { size = 6; num = byte&0x03u; }
    else if( byte <= 0xfdu ) { size = 7; num = byte&0x01u; }
    else if( byte <= 0xfeu ) size = 8;
    else size = 9;
 
    for( int i=0; i<size-1; ++i ){
        unsigned char cur;
        stream >> cur;
        num = (num<<8u)|cur;
    }
    sml.num = num;
    return stream;
}
inline ostream& operator<< ( ostream& stream, const smalluint& sml )
{
    uint64_t num = sml.num;
    unsigned char bytes[9];
    int size;
    if(      num <= 0x000000000000007full ) size = 1;
    else if( num <= 0x0000000000003fffull ) size = 2;
    else if( num <= 0x00000000001fffffull ) size = 3;
    else if( num <= 0x000000000fffffffull ) size = 4;
    else if( num <= 0x00000007ffffffffull ) size = 5;
    else if( num <= 0x000003ffffffffffull ) size = 6;
    else if( num <= 0x0001ffffffffffffull ) size = 7;
    else if( num <= 0x00ffffffffffffffull ) size = 8;
    else size = 9;
 
    for( int j=size-1; j>=0; --j ){
        bytes[j] = (num&0xffull);
        num >>= 8u;
    }
    bytes[0] |= (0xfff00ull>>(size-1));
    stream.write((char*)bytes,size);
 
    return stream;
}
 
namespace {
static const std::size_t MUI_TCP_BUF_SIZE = 4000;
int SYSCHECK(int value){ // carefully use this. just like "SYSCHECK(close(fd));"
    int err = errno;
    if(value == -1) throw std::system_error(err, std::system_category());
    return value;
}
}
 
struct mutex_timeout: std::runtime_error {
    mutex_timeout(): std::runtime_error("tcp error: lock time out. May be dead locked.") {}
};
 
/* unique_fd_
 * RAII class for file-descripter.
 */
class unique_fd_ {
public:
    unique_fd_(int fd=0): fd_(fd) {}
    unique_fd_(const unique_fd_&) = delete;
    unique_fd_(unique_fd_&& rhs): fd_(rhs.fd_) {
        rhs.fd_ = 0;
    }
    unique_fd_& operator=(const unique_fd_&) = delete;
    unique_fd_& operator=(unique_fd_&& rhs){
        checked_close_();
        fd_ = rhs.fd_;
        rhs.fd_ = 0;
        return *this;
    }
    ~unique_fd_(){ checked_close_(); }
    operator int () const { return fd_; }
    void reset(int fd) {
        checked_close_();
        fd_ = fd;
    }
    void swap(unique_fd_& rhs) { std::swap(fd_, rhs.fd_); }
private:
    void checked_close_(){ if( fd_ ) close(fd_); }
    int fd_;
};
 
/* poll_scheduler
 * A wrapper for epoll. It calls callback functions if the fd binded to callback becomes ready.
 * port to select? kqueue?
 */
class poll_scheduler {
public:
    static const uint64_t IN = 1;
    static const uint64_t OUT = 4;
    static const uint64_t CALL = 32;
    static const uint64_t ET = 1u<<31; // edge triggered
    
    typedef std::function<void(uint64_t)> callback_t;
    
    poll_scheduler(int max){
        buf_.reserve(max+1);
        callbacks_.reserve(max+1);
 
        epfd_ = SYSCHECK(epoll_create(max));
 
        int pipes[2];
        SYSCHECK(pipe(pipes));
        pipe_read_.reset(pipes[0]);
        pipe_write_.reset(pipes[1]);
        add(pipe_read_, IN, std::bind(&poll_scheduler::pop_event_, this));
    }
    
    ~poll_scheduler() {
        for( auto& a: callbacks_ )
            SYSCHECK(epoll_ctl(epfd_, EPOLL_CTL_DEL, a.first, NULL));
    }
 
    poll_scheduler(const poll_scheduler&) = delete;
    poll_scheduler& operator=(const poll_scheduler&) = delete;
    
    void run() {
        buf_.resize(callbacks_.size());
        int nfd = epoll_wait(epfd_, buf_.data(), buf_.size(), 100);
        int err = errno;
        if( nfd == 0 ) return;;
        if( nfd == -1 && err == EAGAIN ) return;
        else if( nfd == -1 ) throw std::system_error(err,std::system_category());
        for( int i=0; i<nfd; ++i )
            callbacks_[buf_[i].data.fd](translate_(buf_[i].events));
    }
 
    void add(int fd, uint64_t default_events, callback_t callback) {
        callbacks_.emplace(fd, std::move(callback));
        epoll_event ev;
        ev.data.fd = fd;
        ev.events = translate_from_(default_events) | EPOLLET;
        try {
            SYSCHECK(epoll_ctl(epfd_, EPOLL_CTL_ADD, fd, &ev));
        } catch(...) {
            callbacks_.erase(callbacks_.find(fd));
            throw;
        }
    }
    void schedule(int fd) {
        SYSCHECK(write(pipe_write_, &fd, sizeof(int)));
    }
private:
    void pop_event_() {
        int fd;
        SYSCHECK(read(pipe_read_, &fd, sizeof(int)));
        if( fd != -1 )
            callbacks_[fd](CALL);
    }
    uint64_t translate_(uint64_t ev) {
        return (ev&EPOLLIN?IN:0u) | (ev&EPOLLOUT?OUT:0u) | (ev&EPOLLET?ET:0u);
    }
    uint64_t translate_from_(uint64_t ps_ev) {
        return (ps_ev&IN?EPOLLIN:0u) | (ps_ev&OUT?EPOLLOUT:0u) |(ps_ev&ET?EPOLLET:0u);
    }
    
    std::unordered_map<int, callback_t> callbacks_;
    std::vector<epoll_event> buf_;
    unique_fd_ epfd_;
    unique_fd_ pipe_read_, pipe_write_;
};
 
namespace {
/* read_buffer
 * This manages buffers which will be used by read(3).
 * Internal use only.
 */
struct read_buffer : istream {
    static const int BUFSIZE = MUI_TCP_BUF_SIZE;
    struct node_t { char buf[BUFSIZE]; };
    
    std::pair<int,char*> get_buffer() {
        if(nodes_.empty() || tail_ == BUFSIZE) {
            nodes_.emplace_back();
            tail_ = 0u;
        }
        return std::make_pair(BUFSIZE-tail_, nodes_.back().buf+tail_);
    }
    void move_tail(int size) { // consume bytes. size must be <= get_buffer().first
        psize_ += size;
        size_ += size;
        tail_ += size;
    }
 
    void read(char* dest, std::size_t size) { // read bytes and seek
        size_ -= size;
        while(size) {
            std::size_t sz = std::min<std::size_t>(size, (pcur_!=nodes_.size()-1?BUFSIZE-ccur_:tail_-ccur_));
            dest = std::copy_n(nodes_[pcur_].buf+ccur_, sz, dest);
            size -= sz;
            ccur_ += sz;
            if( size ) {
                ++pcur_;
                ccur_ = 0u;
            }
        }
    }
    void revert() { // undo seeks
        pcur_ = 0u;
        ccur_ = head_;
        size_ = psize_;
    }
    void detach() { // erase the used bytes
        nodes_.erase(nodes_.begin(), nodes_.begin()+pcur_);
        head_ = ccur_;
        psize_ = size_;
    }
    std::size_t size() const { return size_; }
    
    std::deque<node_t> nodes_;
    std::size_t psize_ = 0u;
    std::size_t size_ = 0u;
    std::size_t pcur_ = 0u;
    unsigned ccur_ = 0u;
    unsigned head_ = 0u;
    unsigned tail_ = 0u;
};
}
 
class read_que {
public:
    read_que( unique_fd_&& fd, std::function<void(message)> callback )
    : fd_(std::move(fd)), callback_(std::move(callback)) {}
    read_que( read_que&& ) = default;
    read_que& operator=( read_que&& ) = default;
    
    void try_recv(int){
        while(true){
            std::pair<int,char*> p = buf_.get_buffer();
            ssize_t r = read(fd_, p.second, p.first);
            int err = errno;
            if( r > 0 ){
                buf_.move_tail(r);
                char c;
                buf_ >> c;
                buf_.revert();
                std::uint64_t smlsize = smalluint::detect_size(c);
                if( buf_.size() < smlsize) continue;
 
                smalluint vecsize;
                buf_ >> vecsize;
                buf_.revert();
                if( buf_.size() < smlsize + vecsize.num ) continue;
 
                message msg;
                buf_ >> msg;
                callback_(std::move(msg));
                buf_.detach();
                continue;
            } else if( r == -1 && (err == EAGAIN || err == EWOULDBLOCK) ){
                break;
            }else if( r == -1 ) throw std::system_error(err,std::system_category());
            else throw std::runtime_error("tcp connection is broken.\n");
        }
    }
    int get_fd() const { return fd_; }
private:
    unique_fd_ fd_;
    std::function<void(message)> callback_;
    read_buffer buf_;
};
 
class write_que {
public:
    explicit write_que(unique_fd_&& fd): fd_(std::move(fd)) {}
    write_que(write_que&& rhs): fd_(std::move(rhs.fd_)), send_(std::move(send_)) {}
    write_que& operator=(write_que&& rhs) {
        std::swap(fd_, rhs.fd_);
        send_.swap(rhs.send_);
        return *this;
    }
    
    void try_send(int){
        while(true) {
            mutex_.lock();
            auto iter = send_.begin();
            auto end  = send_.end();
            mutex_.unlock();
            if(iter == end) break;
            auto& p = *iter;
 
            std::size_t sz = p.second.size() - p.first;
            char* buf = p.second.data() + p.first;
            ssize_t r = write(fd_, buf, sz);
            int err = errno;
            if( r == sz ) {
                std::lock_guard<std::mutex> lock(mutex_);
                send_.pop_front();
            } else if( r > 0 ) {
                p.first += r;
            } else if( r == -1 && (err == EAGAIN || err == EWOULDBLOCK) ){
                break;
            } else if( r == -1 ) throw std::system_error(err,std::system_category());
            else throw std::runtime_error("tcp connection is broken.\n");
        }
    }
    void push( std::vector<char> data ){
        if( data.size() > std::numeric_limits<ssize_t>::max() ) 
            throw std::range_error("cannot send data larger than ssize_t");
        if( !data.empty() ) {
            std::lock_guard<std::mutex> lock(mutex_);
            send_.emplace_back(0u, std::move(data));
        }
    }
    int get_fd() const { return fd_; }
private:
    unique_fd_ fd_;
    std::mutex mutex_;
    std::list<std::pair<std::size_t,std::vector<char> > > send_;
};
 
class comm_fd: public communicator {
public:
    comm_fd(int local_rank, int local_size, std::vector<unique_fd_>&& wfds, std::vector<unique_fd_>&& rfds)
    : local_rank_(local_rank), local_size_(local_size), remote_size_(wfds.size()), poll_(rfds.size()+wfds.size()) {
        assert(wfds.size() == rfds.size());
        die_.store(false);
 
        rques_.reserve(rfds.size());
        wques_.reserve(wfds.size());
 
        for(auto& r: rfds) rques_.emplace_back(std::move(r),std::bind(&comm_fd::push_msg_,this,std::placeholders::_1));
        for(auto& q: rques_)
            poll_.add(q.get_fd(), poll_scheduler::IN | poll_scheduler::ET, std::bind(&read_que::try_recv,&q,std::placeholders::_1));
        for(auto& w: wfds) wques_.emplace_back(std::move(w));
        for(auto& q: wques_)
            poll_.add(q.get_fd(), poll_scheduler::OUT | poll_scheduler::ET, std::bind(&write_que::try_send,&q,std::placeholders::_1));
 
        std::thread th = std::thread(std::bind(&comm_fd::run_, this));
        thread_.swap(th);
    }
    
    ~comm_fd() {
        die_.store(true);
        poll_.schedule(-1);
        thread_.join();
    }
    
    int local_rank() const { return local_rank_; }
    int local_size() const { return local_size_; }
    int remote_size() const { return remote_size_; }
protected:
    void send_impl_(message msg, const std::vector<bool>& dest) {
        std::vector<char> v(streamed_size(msg));
        auto stream = make_ostream(v.begin());
        stream << msg;
        for( int i=0; i<remote_size(); ++i ) {
            if(dest[i]){
                wques_[i].push(v);
                poll_.schedule(wques_[i].get_fd());
            }
        }
    }
    message recv_impl_() {
        std::unique_lock<std::mutex> lock(recv_mutex_);
        recv_cv_.wait(lock, [=](){ return !mesgs_.empty(); });
        message msg = std::move(mesgs_.front());
        mesgs_.pop_front();
        return msg;
    }
private:
    void push_msg_(message msg) {
        std::unique_lock<std::mutex> lock(recv_mutex_);
        mesgs_.emplace_back(std::move(msg));
        recv_cv_.notify_one(); // only ONE thrad can cann comm_fd.recv_impl_();
    }
 
    void run_() {
        while(true){
            if( die_ ) break;
            poll_.run();
        }
    }
 
    int local_rank_, local_size_, remote_size_;
 
    std::thread thread_;    
    std::atomic_bool die_;
 
    poll_scheduler poll_;
    std::vector<read_que> rques_;
    std::vector<write_que> wques_;
 
    std::mutex recv_mutex_;
    std::condition_variable recv_cv_;
    std::list<message> mesgs_;
};
 
inline communicator* create_comm_tcp( const char* str ){
    uri u(str);
    bool is_server = u.host().empty();
 
    struct free_addrinfo_ { void operator()( addrinfo* ptr ) const { freeaddrinfo(ptr); } };
    unique_fd_ sock;
    std::unique_ptr<addrinfo,free_addrinfo_> info;
    addrinfo hint, *tmp=NULL;
 
    std::memset((char*)&hint, 0, sizeof(hint));
    hint.ai_family = AF_INET;   // use ip v4
    hint.ai_socktype = SOCK_STREAM; // TCP/IP
    hint.ai_flags = AI_NUMERICSERV|(is_server?AI_PASSIVE:0);
    
    if(int err = getaddrinfo(is_server?0:u.host().c_str(), u.path().c_str(), &hint, &tmp))
        throw std::runtime_error(gai_strerror(err));
    info.reset(tmp);
    
    for(; tmp!=NULL; tmp = tmp->ai_next){
        int s = socket(tmp->ai_family, tmp->ai_socktype, tmp->ai_protocol);
        if(s == -1 ) continue;
        sock.reset(s);
        
        if( is_server ){
            if( bind(sock, tmp->ai_addr, tmp->ai_addrlen) != -1 ) break;
        } else {
            if( connect(sock, tmp->ai_addr, tmp->ai_addrlen) != -1 ) break;
        }
    }
    if(tmp == NULL) throw std::runtime_error("tcp connection error.");
    
    unique_fd_ wfd, rfd;
    
    if( is_server ){
        SYSCHECK(listen(sock,1));
        wfd.reset(SYSCHECK(accept(sock,0,0)));
        SYSCHECK(fcntl(wfd,F_SETFL,O_NONBLOCK));
        rfd.reset(SYSCHECK(dup(wfd)));
    } else {
        SYSCHECK(fcntl(sock,F_SETFL,O_NONBLOCK));
        wfd.swap(sock);
        rfd.reset(SYSCHECK(dup(wfd)));
    }
    std::vector<mui::unique_fd_> w; w.emplace_back(std::move(wfd));
    std::vector<mui::unique_fd_> r; r.emplace_back(std::move(rfd));
    return static_cast<communicator*>(new comm_fd(0,1,std::move(w),std::move(r)));
}
 
const static bool comm_tcp_registered_ = comm_factory::instance().link( "tcp", create_comm_tcp );
 
inline communicator* create_comm_shm( const char* str ){
    uri u(str);
    std::string patha = std::string("/dev/shm/") + u.host() + "_A";
    std::string pathb = std::string("/dev/shm/") + u.host() + "_B";
    bool is_wr = false;
    
    if(mkfifo(patha.c_str(), S_IRUSR | S_IWUSR)) {
        int err = errno;
        if(err != EEXIST) throw std::system_error(err, std::system_category());
        is_wr = true;
    }
    if(!is_wr)
        if(int ret = mkfifo(pathb.c_str(), S_IRUSR | S_IWUSR))
            throw std::system_error(errno, std::system_category());
    
    unique_fd_ afd = SYSCHECK(open(patha.c_str(), is_wr?O_RDONLY:O_WRONLY));
    SYSCHECK(fcntl(afd,F_SETFL,O_NONBLOCK));
    unique_fd_ bfd = SYSCHECK(open(pathb.c_str(), is_wr?O_WRONLY:O_RDONLY));
    SYSCHECK(fcntl(bfd,F_SETFL,O_NONBLOCK));
    
    std::vector<mui::unique_fd_> w; w.emplace_back(std::move(is_wr?bfd:afd));
    std::vector<mui::unique_fd_> r; r.emplace_back(std::move(is_wr?afd:bfd));
    return static_cast<communicator*>(new comm_fd(0,1,std::move(w),std::move(r)));
}
 
const static bool comm_shm_registered_ = comm_factory::instance().link( "shm", create_comm_shm);
 
}
 
#if 0
void print( mui::message msg )
{
    auto stream = mui::make_istream(msg.data());
    std::string str;
    stream >> str;
    std::cerr<< "GET MESSAGE!!\t" << str << std::endl;
}
 
int main(int argc, char **argv)
{
    //std::unique_ptr<mui::communicator> c(mui::create_comm_tcp(argc<2?"tcp:///37129":"tcp://localhost/37129"));
    std::unique_ptr<mui::communicator> c(mui::create_comm_shm("shm://my_pipe/"));
 
    c->send(mui::message::make("test",std::string("Hello, ")));
    mui::message msg = c->recv();
    print(msg);
    sleep(3);
    return 0;
}
#endif
 
#endif