1 /** 2 * Copyright 2021 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_DISTRIBUTED_RPC_TCP_CONNECTION_POOL_H_ 18 #define MINDSPORE_CCSRC_DISTRIBUTED_RPC_TCP_CONNECTION_POOL_H_ 19 20 #include <map> 21 #include <set> 22 #include <string> 23 #include <mutex> 24 25 #include "include/backend/distributed/rpc/tcp/constants.h" 26 #include "distributed/rpc/tcp/connection.h" 27 28 namespace mindspore { 29 namespace distributed { 30 namespace rpc { 31 struct ConnectionInfo { 32 int socket_fd; 33 std::string from; 34 std::string to; 35 DeleteCallBack delete_callback; 36 }; 37 38 /* 39 * Maintains a collection of reusable connections. 40 */ 41 class ConnectionPool { 42 public: ConnectionPool()43 ConnectionPool() : double_link_(false) {} 44 ~ConnectionPool() = default; 45 46 void Finalize(); 47 48 /* 49 * Operations for ConnectionInfo. 50 */ 51 void AddConnInfo(int socket_fd, const std::string &dst_url, DeleteCallBack delcb); 52 bool ReverseConnInfo(int from_socket_fd, int to_socket_fd); 53 54 /* 55 * Operations for Connection. 56 */ 57 // Add a connection. 58 void AddConnection(Connection *conn); 59 60 // Find connection. 61 Connection *FindConnection(const std::string &dst_url); 62 63 // Delete connection. 64 void DeleteConnection(const std::string &dst_url); 65 void DeleteAllConnections(std::map<std::string, Connection *> *alllinks) const; 66 67 // Close connection. 68 void CloseConnection(Connection *conn); 69 70 // Single link or double link. 71 void SetLinkPattern(bool linkPattern); 72 73 void ResetAllConnMetrics(); 74 75 private: 76 ConnectionInfo *FindConnInfo(int socket_fd, const std::string &dst_url); 77 78 void DeleteConnInfo(int socket_fd); 79 void DeleteConnInfo(Connection *conn); 80 void DeleteAllConnInfos(); 81 82 bool double_link_; 83 84 // to_url=tcp@ip:port, event struct 85 std::map<std::string, Connection *> local_conns_; 86 87 // Maintains the remote connections by remote server addresses. 88 std::map<std::string, Connection *> remote_conns_; 89 90 // Maintains the connections by remote server addresses. 91 std::map<std::string, Connection *> connections_; 92 93 // each to_url has two fds at most, and each fd has multiple linkinfos 94 std::map<int, std::set<ConnectionInfo *>> conn_infos_; 95 96 // This mutex is used for protecting the modification of connections. 97 std::mutex mutex_; 98 99 friend class TCPComm; 100 }; 101 } // namespace rpc 102 } // namespace distributed 103 } // namespace mindspore 104 #endif 105