socket.h

/************************************************************
*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements.  See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership.  The ASF licenses this file
* to you 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.
*
*************************************************************/

#ifndef SINGA_COMMUNICATION_SOCKET_H_
#define SINGA_COMMUNICATION_SOCKET_H_

#ifdef USE_ZMQ
#include <czmq.h>
#endif
#include <map>
#include <string>
#include <vector>
#include "communication/msg.h"

namespace singa {

const std::string kInprocRouterEndpoint = "inproc://router";

class SocketInterface {
 public:
  virtual ~SocketInterface() {}
  virtual int Send(Msg** msg) = 0;
  virtual Msg* Receive() = 0;
  virtual void* InternalID() const = 0;
};

class Poller {
 public:
  Poller();
  explicit Poller(SocketInterface* socket);
  void Add(SocketInterface* socket);
  SocketInterface* Wait(int duration);

  virtual bool Terminated();

 protected:
#ifdef USE_ZMQ
  zpoller_t *poller_;
  std::map<zsock_t*, SocketInterface*> zsock2Socket_;
#endif
};

class Dealer : public SocketInterface {
 public:
  /*
   * @param id Local dealer ID within a procs if the dealer is from worker or
   * server thread, starts from 1 (0 is used by the router); or the connected
   * remote procs ID for inter-process dealers from the stub thread.
   */
  Dealer();
  explicit Dealer(int id);
  ~Dealer() override;
  int Connect(const std::string& endpoint);
  int Send(Msg** msg) override;
  Msg* Receive() override;
  void* InternalID() const override;

 protected:
  int id_ = -1;
#ifdef USE_ZMQ
  zsock_t* dealer_ = nullptr;
  zpoller_t* poller_ = nullptr;
#endif
};

class Router : public SocketInterface {
 public:
  Router();
  explicit Router(int bufsize);
  ~Router() override;
  int Bind(const std::string& endpoint);
  int Send(Msg** msg) override;
  Msg* Receive() override;
  void* InternalID() const override;

 protected:
  int nBufmsg_ = 0;
  int bufsize_ = 100;
#ifdef USE_ZMQ
  zsock_t* router_ = nullptr;
  zpoller_t* poller_ = nullptr;
  std::map<int, zframe_t*> id2addr_;
  std::map<int, std::vector<zmsg_t*>> bufmsg_;
#endif
};

#ifdef USE_MPI
// TODO(wangsheng): add intra-process communication using shared queue
std::vector<SafeQueue*> MPIQueues;
#endif

}  // namespace singa

#endif  // SINGA_COMMUNICATION_SOCKET_H_