github.com/futurehomeno/fimpgo@v1.14.0/mqtt_conn_pool.go (about) 1 package fimpgo 2 3 import ( 4 "errors" 5 "fmt" 6 log "github.com/sirupsen/logrus" 7 "math/rand" 8 "sync" 9 "time" 10 ) 11 12 type connection struct { 13 mqConnection *MqttTransport 14 isIdle bool 15 startedAt time.Time 16 idleSince time.Time 17 } 18 19 // Connection pool starts at initSize connections and can grow up to maxSize , the pool can shrink back to size defined in "size" variable 20 21 type MqttConnectionPool struct { 22 mux sync.RWMutex 23 connTemplate MqttConnectionConfigs 24 connPool map[int]*connection 25 clientIdPrefix string 26 initSize int // init size 27 size int // normal size 28 maxSize int // max size 29 maxIdleAge time.Duration // Defines how long idle connection can stay in the pool before it gets destroyed 30 poolCheckTick *time.Ticker 31 isActive bool 32 } 33 34 func NewMqttConnectionPool(initSize, size, maxSize int, maxAge time.Duration, connTemplate MqttConnectionConfigs, clientIdPrefix string) *MqttConnectionPool { 35 if maxAge == 0 { 36 maxAge = 30 * time.Second // age in seconds 37 } 38 if maxSize == 0 { 39 maxSize = 20 40 } 41 if size == 0 { 42 size = 2 43 } 44 if initSize > maxSize { 45 initSize = 0 46 } 47 pool := &MqttConnectionPool{connPool: make(map[int]*connection), connTemplate: connTemplate, clientIdPrefix: clientIdPrefix, initSize: initSize, size: size, maxSize: maxSize, maxIdleAge: maxAge} 48 pool.Start() 49 return pool 50 } 51 52 func (cp *MqttConnectionPool) Start() { 53 if !cp.isActive { 54 cp.isActive = true 55 cp.poolCheckTick = time.NewTicker(10 * time.Second) 56 go cp.cleanupProcess() 57 } 58 } 59 60 func (cp *MqttConnectionPool) Stop() { 61 cp.isActive = false 62 } 63 64 func (cp *MqttConnectionPool) TotalConnections() int { 65 cp.mux.RLock() 66 size := len(cp.connPool) 67 cp.mux.RUnlock() 68 return size 69 } 70 71 func (cp *MqttConnectionPool) IdleConnections() int { 72 var size int 73 cp.mux.RLock() 74 for i := range cp.connPool { 75 if cp.connPool[i].isIdle { 76 size++ 77 } 78 } 79 cp.mux.RUnlock() 80 return size 81 } 82 83 func (cp *MqttConnectionPool) createConnection() (int, error) { 84 connId := cp.genConnId() 85 if len(cp.connPool) >= cp.maxSize { 86 log.Error("<mq-pool> Too many connections") 87 return 0, errors.New("too many connections") 88 } 89 conf := cp.connTemplate 90 conf.ClientID = fmt.Sprintf("%s_%d", cp.clientIdPrefix, connId) 91 newConnection := NewMqttTransportFromConfigs(conf) 92 err := newConnection.Start() 93 cp.connPool[connId] = &connection{ 94 mqConnection: newConnection, 95 isIdle: false, 96 startedAt: time.Now(), 97 } 98 log.Debugf("New connection %d created . Pool size = %d", connId, len(cp.connPool)) 99 return connId, err 100 } 101 102 // BorrowConnection returns first available connection from the pool or creates new connection 103 func (cp *MqttConnectionPool) BorrowConnection() (int, *MqttTransport, error) { 104 defer cp.mux.Unlock() 105 cp.mux.Lock() 106 for i := range cp.connPool { 107 if cp.connPool[i].isIdle { 108 if cp.connPool[i].mqConnection.Client().IsConnected() { 109 cp.connPool[i].isIdle = false 110 return i, cp.connPool[i].mqConnection, nil 111 } else { 112 break 113 } 114 } 115 } 116 connId, err := cp.createConnection() 117 return connId, cp.getConnectionById(connId), err 118 } 119 120 // ReturnConnection returns connection to pool by setting inUse status to false 121 func (cp *MqttConnectionPool) ReturnConnection(connId int) { 122 defer cp.mux.RUnlock() 123 cp.mux.RLock() 124 con, ok := cp.connPool[connId] 125 if ok { 126 con.mqConnection.UnsubscribeAll() 127 con.isIdle = true 128 con.idleSince = time.Now() 129 log.Debugf("Connection %d returned to pool.", connId) 130 } 131 } 132 133 // getConnectionById returns connection from pool or creates new connection 134 func (cp *MqttConnectionPool) getConnectionById(connId int) *MqttTransport { 135 conn, ok := cp.connPool[connId] 136 if ok { 137 return conn.mqConnection 138 } 139 return nil 140 } 141 142 func (cp *MqttConnectionPool) genConnId() int { 143 rand.Seed(int64(time.Now().Nanosecond())) 144 for { 145 id := rand.Int() 146 if _, ok := cp.connPool[id]; !ok { 147 return id 148 } 149 } 150 } 151 152 func (cp *MqttConnectionPool) cleanupProcess() { 153 for { 154 <-cp.poolCheckTick.C 155 if !cp.isActive { 156 break 157 } 158 cp.mux.Lock() 159 if len(cp.connPool) > cp.size { 160 for i := range cp.connPool { 161 if cp.connPool[i].isIdle { 162 if (time.Since(cp.connPool[i].idleSince) > (cp.maxIdleAge)) && (len(cp.connPool) > cp.size) { 163 log.Debugf("<conn-pool> Destroying old connection") 164 conn := cp.getConnectionById(i) 165 conn.Stop() 166 delete(cp.connPool, i) // it is safe to delete map element in the loop 167 } else { 168 //log.Debugf("<conn-pool> Nothing to clean") 169 } 170 171 } 172 } 173 } 174 cp.mux.Unlock() 175 } 176 }