github.com/pawelgaczynski/gain@v0.4.0-alpha.0.20230821120126-41f1e60a18da/conn_manager.go (about) 1 // Copyright (c) 2023 Paweł Gaczyński 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package gain 16 17 import "sync/atomic" 18 19 type connectionManager struct { 20 connections map[int]*connection 21 connectionsCount atomic.Int32 22 closing bool 23 releaseFdSet map[int]void 24 keyPool *keyPool 25 } 26 27 func (c *connectionManager) fork(conn *connection, write bool) *connection { 28 key := int(c.keyPool.get()) 29 30 forked := getConnection() 31 forked = conn.fork(forked, key, write) 32 c.connections[key] = forked 33 c.connectionsCount.Add(1) 34 35 return forked 36 } 37 38 func (c *connectionManager) getFd(fd int) *connection { 39 return c.get(fd, fd) 40 } 41 42 func (c *connectionManager) get(key int, fd int) *connection { 43 conn, ok := c.connections[key] 44 if ok { 45 return conn 46 } 47 48 conn = getConnection() 49 conn.fd = fd 50 conn.key = key 51 c.connections[key] = conn 52 c.connectionsCount.Add(1) 53 54 return conn 55 } 56 57 func (c *connectionManager) release(key int) { 58 conn, ok := c.connections[key] 59 if ok { 60 delete(c.connections, key) 61 c.connectionsCount.Add(-1) 62 putConnection(conn) 63 delete(c.releaseFdSet, key) 64 c.keyPool.put(uint64(conn.key)) 65 } 66 } 67 68 func (c *connectionManager) close(callback func(conn *connection) bool, fdSkipped int) { 69 c.closing = true 70 c.releaseFdSet = make(map[int]void) 71 72 for _, value := range c.connections { 73 if value.fd == fdSkipped { 74 continue 75 } 76 77 if callback(value) { 78 c.releaseFdSet[value.fd] = member 79 } 80 } 81 } 82 83 func (c *connectionManager) allClosed() bool { 84 return c.closing && len(c.releaseFdSet) == 0 85 } 86 87 func (c *connectionManager) activeConnections() int { 88 return int(c.connectionsCount.Load()) 89 } 90 91 func newConnectionManager() *connectionManager { 92 return &connectionManager{ 93 connections: make(map[int]*connection), 94 keyPool: newKeyPool(), 95 } 96 }