github.com/anycable/anycable-go@v1.5.1/node/disconnect_queue.go (about)

     1  package node
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"log/slog"
     7  	"sync"
     8  	"time"
     9  )
    10  
    11  // DisconnectQueueConfig contains DisconnectQueue configuration
    12  type DisconnectQueueConfig struct {
    13  	// Limit the number of Disconnect RPC calls per second
    14  	Rate int
    15  	// The size of the channel's buffer for disconnect requests
    16  	Backlog int
    17  	// How much time wait to call all enqueued calls at exit (in seconds) [DEPREACTED]
    18  	ShutdownTimeout int
    19  }
    20  
    21  // NewDisconnectQueueConfig builds a new config
    22  func NewDisconnectQueueConfig() DisconnectQueueConfig {
    23  	return DisconnectQueueConfig{Rate: 100, Backlog: 4096}
    24  }
    25  
    26  // DisconnectQueue is a rate-limited executor
    27  type DisconnectQueue struct {
    28  	node *Node
    29  	// Throttling rate
    30  	rate time.Duration
    31  	// Call RPC Disconnect for connections
    32  	disconnect chan *Session
    33  	// Logger with context
    34  	log *slog.Logger
    35  	// Control channel to shutdown the executer
    36  	shutdown chan struct{}
    37  	// Executer stopped status
    38  	isStopped bool
    39  	// Mutex to work with stopped status concurrently
    40  	mu sync.Mutex
    41  }
    42  
    43  // NewDisconnectQueue builds new queue with a specified rate (max calls per second)
    44  func NewDisconnectQueue(node *Node, config *DisconnectQueueConfig, l *slog.Logger) *DisconnectQueue {
    45  	rateDuration := time.Millisecond * time.Duration(1000/config.Rate)
    46  
    47  	ctx := l.With("context", "disconnector")
    48  
    49  	ctx.Debug("calls rate", "rate", rateDuration)
    50  
    51  	return &DisconnectQueue{
    52  		node:       node,
    53  		disconnect: make(chan *Session, config.Backlog),
    54  		rate:       rateDuration,
    55  		log:        ctx,
    56  		shutdown:   make(chan struct{}, 1),
    57  	}
    58  }
    59  
    60  // Run starts queue
    61  func (d *DisconnectQueue) Run() error {
    62  	throttle := time.NewTicker(d.rate)
    63  	defer throttle.Stop()
    64  
    65  	for {
    66  		select {
    67  		case session := <-d.disconnect:
    68  			<-throttle.C
    69  			d.node.DisconnectNow(session) //nolint:errcheck
    70  		case <-d.shutdown:
    71  			return nil
    72  		}
    73  	}
    74  }
    75  
    76  // Shutdown stops throttling and makes requests one by one
    77  func (d *DisconnectQueue) Shutdown(ctx context.Context) error {
    78  	d.mu.Lock()
    79  	if d.isStopped {
    80  		d.mu.Unlock()
    81  		return nil
    82  	}
    83  
    84  	d.isStopped = true
    85  	d.shutdown <- struct{}{}
    86  	d.mu.Unlock()
    87  
    88  	left := len(d.disconnect)
    89  	actual := 0
    90  
    91  	if left == 0 {
    92  		return nil
    93  	}
    94  
    95  	defer func() {
    96  		d.log.Info("disconnected sessions", "num", actual)
    97  	}()
    98  
    99  	deadline, ok := ctx.Deadline()
   100  
   101  	if ok {
   102  		timeLeft := time.Until(deadline)
   103  
   104  		d.log.Info("invoking remaining disconnects", "interval", timeLeft.Seconds(), "num", left)
   105  	} else {
   106  		d.log.Info("invoking remaining disconnects", "num", left)
   107  	}
   108  
   109  	for {
   110  		select {
   111  		case session := <-d.disconnect:
   112  			d.node.DisconnectNow(session) // nolint:errcheck
   113  
   114  			actual++
   115  		case <-ctx.Done():
   116  			return fmt.Errorf("had no time to invoke Disconnect calls: ~%d", len(d.disconnect))
   117  		default:
   118  			return nil
   119  		}
   120  	}
   121  }
   122  
   123  // Enqueue adds session to the disconnect queue
   124  func (d *DisconnectQueue) Enqueue(s *Session) error {
   125  	d.mu.Lock()
   126  
   127  	// Check that we're not closed
   128  	if d.isStopped {
   129  		d.mu.Unlock()
   130  		return nil
   131  	}
   132  
   133  	d.mu.Unlock()
   134  
   135  	d.disconnect <- s
   136  
   137  	return nil
   138  }
   139  
   140  // Size returns the number of enqueued tasks
   141  func (d *DisconnectQueue) Size() int {
   142  	return len(d.disconnect)
   143  }