github.com/e154/smart-home@v0.17.2-0.20240311175135-e530a6e5cd45/system/gate/server/wsp/pools.go (about)

     1  // This file is part of the Smart Home
     2  // Program complex distribution https://github.com/e154/smart-home
     3  // Copyright (C) 2024, Filippov Alex
     4  //
     5  // This library is free software: you can redistribute it and/or
     6  // modify it under the terms of the GNU Lesser General Public
     7  // License as published by the Free Software Foundation; either
     8  // version 3 of the License, or (at your option) any later version.
     9  //
    10  // This library is distributed in the hope that it will be useful,
    11  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    12  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
    13  // Library General Public License for more details.
    14  //
    15  // You should have received a copy of the GNU Lesser General Public
    16  // License along with this library.  If not, see
    17  // <https://www.gnu.org/licenses/>.
    18  
    19  package wsp
    20  
    21  import (
    22  	"fmt"
    23  	"strconv"
    24  	"strings"
    25  	"sync"
    26  	"time"
    27  
    28  	"github.com/gorilla/websocket"
    29  )
    30  
    31  type Pools struct {
    32  	timeout     time.Duration
    33  	idleTimeout time.Duration
    34  	*sync.Mutex
    35  	pools map[PoolID]*Pool
    36  }
    37  
    38  func NewPools(timeout, idleTimeout time.Duration) *Pools {
    39  	return &Pools{
    40  		timeout:     timeout,
    41  		idleTimeout: idleTimeout,
    42  		Mutex:       &sync.Mutex{},
    43  		pools:       make(map[PoolID]*Pool),
    44  	}
    45  }
    46  
    47  func (p *Pools) Shutdown() {
    48  	for _, pool := range p.pools {
    49  		pool.Shutdown()
    50  	}
    51  }
    52  
    53  func (p *Pools) RegisterConnection(ws *websocket.Conn) (err error) {
    54  
    55  	// 2. Wait a greeting message from the peer and parse it
    56  	// The first message should contains the remote Proxy name and size
    57  	_, greeting, err := ws.ReadMessage()
    58  	if err != nil {
    59  		err = fmt.Errorf("Unable to read greeting message : %s", err)
    60  		return
    61  	}
    62  
    63  	// Parse the greeting message
    64  	split := strings.Split(string(greeting), "_")
    65  	poolID := PoolID(split[0])
    66  	size, err := strconv.Atoi(split[1])
    67  	if err != nil {
    68  		err = fmt.Errorf("Unable to parse greeting message : %s", err)
    69  		return
    70  	}
    71  
    72  	p.Lock()
    73  	defer p.Unlock()
    74  
    75  	if _, ok := p.pools[poolID]; !ok {
    76  		p.pools[poolID] = NewPool(p.timeout, p.idleTimeout, poolID)
    77  	}
    78  
    79  	// update pool size
    80  	p.pools[poolID].SetSize(size)
    81  
    82  	// Add the WebSocket connection to the pool
    83  	p.pools[poolID].RegisterConnection(ws)
    84  
    85  	return
    86  }
    87  
    88  func (p *Pools) Clean() {
    89  	p.Lock()
    90  	defer p.Unlock()
    91  
    92  	if len(p.pools) == 0 {
    93  		return
    94  	}
    95  
    96  	idle := 0
    97  	busy := 0
    98  	closed := 0
    99  
   100  	for _, pool := range p.pools {
   101  		if pool.IsEmpty() {
   102  			log.Infof("Removing empty connection pool : %p", pool.id)
   103  			pool.Shutdown()
   104  			delete(p.pools, pool.id)
   105  		}
   106  
   107  		ps := pool.Size()
   108  		idle += ps.Idle
   109  		busy += ps.Busy
   110  		closed = ps.Closed
   111  	}
   112  
   113  	log.Infof("%d pools, %d idle, %d busy, %d closed", len(p.pools), idle, busy, closed)
   114  }
   115  
   116  func (p *Pools) IsEmpty() bool {
   117  	p.Lock()
   118  	defer p.Unlock()
   119  	return len(p.pools) == 0
   120  }
   121  
   122  func (p *Pools) MoreThenOne() bool {
   123  	p.Lock()
   124  	defer p.Unlock()
   125  	return len(p.pools) > 1
   126  }
   127  
   128  func (p *Pools) GetPool(id PoolID) (pool *Pool, ok bool) {
   129  	p.Lock()
   130  	defer p.Unlock()
   131  	pool, ok = p.pools[id]
   132  	return
   133  }