1 /** 2 * Copyright 2020 Huawei Technologies Co., Ltd 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 #ifndef MINDSPORE_CCSRC_PS_CORE_COMMUNICATOR_TCP_SERVER_H_ 18 #define MINDSPORE_CCSRC_PS_CORE_COMMUNICATOR_TCP_SERVER_H_ 19 20 #include <event2/buffer.h> 21 #include <event2/bufferevent.h> 22 #include <event2/event.h> 23 #include <event2/listener.h> 24 #include <event2/thread.h> 25 #include <event2/bufferevent_ssl.h> 26 27 #include <exception> 28 #include <functional> 29 #include <iostream> 30 #include <map> 31 #include <memory> 32 #include <mutex> 33 #include <string> 34 #include <vector> 35 #include <thread> 36 #include <atomic> 37 38 #include "ps/core/communicator/tcp_message_handler.h" 39 #include "ps/core/communicator/ssl_wrapper.h" 40 #include "ps/core/cluster_config.h" 41 #include "utils/convert_utils_base.h" 42 #include "ps/core/comm_util.h" 43 #include "ps/constants.h" 44 #include "ps/ps_context.h" 45 #include "ps/core/file_configuration.h" 46 47 namespace mindspore { 48 namespace ps { 49 namespace core { 50 class TcpServer; 51 class TcpConnection { 52 public: TcpConnection(struct bufferevent * bev,const evutil_socket_t & fd,TcpServer * server)53 explicit TcpConnection(struct bufferevent *bev, const evutil_socket_t &fd, TcpServer *server) 54 : buffer_event_(bev), fd_(fd), server_(server) {} 55 TcpConnection(const TcpConnection &); 56 virtual ~TcpConnection(); 57 58 using Callback = std::function<void(const std::shared_ptr<CommMessage>)>; 59 60 void InitConnection(const messageReceive &callback); 61 void SendMessage(const void *buffer, size_t num) const; 62 bool SendMessage(const std::shared_ptr<CommMessage> &message) const; 63 bool SendMessage(const std::shared_ptr<MessageMeta> &meta, const Protos &protos, const void *data, size_t size) const; 64 void OnReadHandler(const void *buffer, size_t numBytes); 65 const TcpServer *GetServer() const; 66 const evutil_socket_t &GetFd() const; 67 void set_callback(const Callback &callback); 68 69 protected: 70 struct bufferevent *buffer_event_; 71 evutil_socket_t fd_; 72 TcpServer *server_; 73 TcpMessageHandler tcp_message_handler_; 74 Callback callback_; 75 }; 76 77 using OnServerReceiveMessage = 78 std::function<void(const std::shared_ptr<TcpConnection> &conn, const std::shared_ptr<MessageMeta> &meta, 79 const Protos &protos, const void *data, size_t size)>; 80 81 class TcpServer { 82 public: 83 using OnConnected = std::function<void(const TcpServer &, const TcpConnection &)>; 84 using OnDisconnected = std::function<void(const TcpServer &, const TcpConnection &)>; 85 using OnAccepted = std::function<std::shared_ptr<TcpConnection>(const TcpServer &)>; 86 using OnTimerOnce = std::function<void(const TcpServer &)>; 87 using OnTimer = std::function<void()>; 88 89 TcpServer(const std::string &address, std::uint16_t port, Configuration *const config); 90 TcpServer(const TcpServer &server); 91 virtual ~TcpServer(); 92 93 void SetServerCallback(const OnConnected &client_conn, const OnDisconnected &client_disconn, 94 const OnAccepted &client_accept); 95 void set_timer_once_callback(const OnTimerOnce &timer); 96 void set_timer_callback(const OnTimer &timer); 97 void Init(); 98 void Start(); 99 void StartWithNoBlock(); 100 void Stop(); 101 void SendToAllClients(const char *data, size_t len); 102 void AddConnection(const evutil_socket_t &fd, std::shared_ptr<TcpConnection> connection); 103 void RemoveConnection(const evutil_socket_t &fd); 104 std::shared_ptr<TcpConnection> GetConnectionByFd(const evutil_socket_t &fd); 105 OnServerReceiveMessage GetServerReceive() const; 106 void SetMessageCallback(const OnServerReceiveMessage &cb); 107 bool SendMessage(const std::shared_ptr<TcpConnection> &conn, const std::shared_ptr<CommMessage> &message); 108 bool SendMessage(const std::shared_ptr<TcpConnection> &conn, const std::shared_ptr<MessageMeta> &meta, 109 const Protos &protos, const void *data, size_t sizee); 110 void SendMessage(const std::shared_ptr<CommMessage> &message); 111 uint16_t BoundPort() const; 112 std::string BoundIp() const; 113 int ConnectionNum() const; 114 const std::map<evutil_socket_t, std::shared_ptr<TcpConnection>> &Connections() const; 115 116 protected: 117 static void ListenerCallback(struct evconnlistener *listener, evutil_socket_t socket, struct sockaddr *saddr, 118 int socklen, void *server); 119 static void SignalCallback(evutil_socket_t sig, std::int16_t events, void *server); 120 static void ReadCallback(struct bufferevent *, void *connection); 121 static void EventCallback(struct bufferevent *, std::int16_t events, void *server); 122 static void TimerCallback(evutil_socket_t fd, int16_t event, void *arg); 123 static void TimerOnceCallback(evutil_socket_t fd, int16_t event, void *arg); 124 static void SetTcpNoDelay(const evutil_socket_t &fd); 125 std::shared_ptr<TcpConnection> onCreateConnection(struct bufferevent *bev, const evutil_socket_t &fd); 126 127 struct event_base *base_; 128 struct event *signal_event_; 129 struct evconnlistener *listener_; 130 std::string server_address_; 131 std::uint16_t server_port_; 132 std::atomic<bool> is_stop_; 133 134 std::map<evutil_socket_t, std::shared_ptr<TcpConnection>> connections_; 135 OnConnected client_connection_; 136 OnDisconnected client_disconnection_; 137 OnAccepted client_accept_; 138 std::mutex connection_mutex_; 139 OnServerReceiveMessage message_callback_; 140 OnTimerOnce on_timer_once_callback_; 141 OnTimer on_timer_callback_; 142 // The Configuration file 143 Configuration *config_; 144 int64_t max_connection_; 145 }; 146 } // namespace core 147 } // namespace ps 148 } // namespace mindspore 149 #endif // MINDSPORE_CCSRC_PS_CORE_COMMUNICATOR_TCP_SERVER_H_ 150