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  }