tcp_socket.h

/*!
 *  Copyright (c) 2019 by Contributors
 * \file tcp_socket.h
 * \brief TCP socket for DGL distributed training.
 */
#ifndef DGL_RPC_NETWORK_TCP_SOCKET_H_
#define DGL_RPC_NETWORK_TCP_SOCKET_H_

#ifdef _WIN32
#include <winsock2.h>
#include <ws2tcpip.h>

#pragma comment(lib, "Ws2_32.lib")
#else  // !_WIN32
#include <sys/socket.h>
#endif  // _WIN32
#include <string>

namespace dgl {
namespace network {

/*!
 * \brief TCPSocket is a simple wrapper around a socket.
 * It supports only TCP connections.
 */
class TCPSocket {
 public:
  /*!
   * \brief TCPSocket constructor
   */
  TCPSocket();

  /*!
   * \brief TCPSocket deconstructor
   */
  ~TCPSocket();

  /*!
   * \brief Connect to a given server address
   * \param ip ip address
   * \param port end port
   * \return true for success and false for failure
   */
  bool Connect(const char* ip, int port);

  /*!
   * \brief Bind on the given IP and PORT
   * \param ip ip address
   * \param port end port
   * \return true for success and false for failure
   */
  bool Bind(const char* ip, int port);

  /*!
   * \brief listen for remote connection
   * \param max_connection maximal connection
   * \return true for success and false for failure
   */
  bool Listen(int max_connection);

  /*!
   * \brief wait doe a new connection
   * \param socket new SOCKET will be stored to socket
   * \param ip_client new IP will be stored to ip_client
   * \param port_client new PORT will be stored to port_client
   * \return true for success and false for failure
   */
  bool Accept(TCPSocket* socket, std::string* ip_client, int* port_client);

  /*!
   * \brief SetNonBlocking() is needed refering to this example of epoll:
   * http://www.kernel.org/doc/man-pages/online/pages/man4/epoll.4.html
   * \param flag true for nonblocking, false for blocking
   * \return true for success and false for failure
   */
  bool SetNonBlocking(bool flag);

  /*!
   * \brief Set timeout for socket
   * \param timeout seconds timeout
   */
  void SetTimeout(int timeout);

  /*!
   * \brief Shut down one or both halves of the connection.
   * \param ways ways for shutdown
   * If ways is SHUT_RD, further receives are disallowed.
   * If ways is SHUT_WR, further sends are disallowed.
   * If ways is SHUT_RDWR, further sends and receives are disallowed.
   * \return true for success and false for failure
   */
  bool ShutDown(int ways);

  /*!
   * \brief close socket.
   */
  void Close();

  /*!
   * \brief Send data.
   * \param data data for sending
   * \param len_data length of data
   * \return return number of bytes sent if OK, -1 on error
   */
  int64_t Send(const char* data, int64_t len_data);

  /*!
   * \brief Receive data.
   * \param buffer buffer for receving
   * \param size_buffer size of buffer
   * \return return number of bytes received if OK, -1 on error
   */
  int64_t Receive(char* buffer, int64_t size_buffer);

  /*!
   * \brief Get socket's file descriptor
   * \return socket's file descriptor
   */
  int Socket() const;

 private:
  /*!
   * \brief socket's file descriptor
   */
  int socket_;
};

}  // namespace network
}  // namespace dgl

#endif  // DGL_RPC_NETWORK_TCP_SOCKET_H_