github.com/MetalBlockchain/metalgo@v1.11.9/pubsub/server.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package pubsub
     5  
     6  import (
     7  	"net/http"
     8  	"sync"
     9  	"time"
    10  
    11  	"github.com/gorilla/websocket"
    12  	"go.uber.org/zap"
    13  
    14  	"github.com/MetalBlockchain/metalgo/utils/logging"
    15  	"github.com/MetalBlockchain/metalgo/utils/set"
    16  	"github.com/MetalBlockchain/metalgo/utils/units"
    17  )
    18  
    19  const (
    20  	// Size of the ws read buffer
    21  	readBufferSize = units.KiB
    22  
    23  	// Size of the ws write buffer
    24  	writeBufferSize = units.KiB
    25  
    26  	// Time allowed to write a message to the peer.
    27  	writeWait = 10 * time.Second
    28  
    29  	// Time allowed to read the next pong message from the peer.
    30  	pongWait = 60 * time.Second
    31  
    32  	// Send pings to peer with this period. Must be less than pongWait.
    33  	pingPeriod = (pongWait * 9) / 10
    34  
    35  	// Maximum message size allowed from peer.
    36  	maxMessageSize = 10 * units.KiB // bytes
    37  
    38  	// Maximum number of pending messages to send to a peer.
    39  	maxPendingMessages = 1024 // messages
    40  
    41  	// MaxBytes the max number of bytes for a filter
    42  	MaxBytes = 1 * units.MiB
    43  
    44  	// MaxAddresses the max number of addresses allowed
    45  	MaxAddresses = 10000
    46  )
    47  
    48  type errorMsg struct {
    49  	Error string `json:"error"`
    50  }
    51  
    52  var upgrader = websocket.Upgrader{
    53  	ReadBufferSize:  readBufferSize,
    54  	WriteBufferSize: writeBufferSize,
    55  	CheckOrigin: func(*http.Request) bool {
    56  		return true
    57  	},
    58  }
    59  
    60  // Server maintains the set of active clients and sends messages to the clients.
    61  type Server struct {
    62  	log  logging.Logger
    63  	lock sync.RWMutex
    64  	// conns a list of all our connections
    65  	conns set.Set[*connection]
    66  	// subscribedConnections the connections that have activated subscriptions
    67  	subscribedConnections *connections
    68  }
    69  
    70  // Deprecated: The pubsub server is deprecated.
    71  func New(log logging.Logger) *Server {
    72  	return &Server{
    73  		log:                   log,
    74  		subscribedConnections: newConnections(),
    75  	}
    76  }
    77  
    78  func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    79  	wsConn, err := upgrader.Upgrade(w, r, nil)
    80  	if err != nil {
    81  		s.log.Debug("failed to upgrade",
    82  			zap.Error(err),
    83  		)
    84  		return
    85  	}
    86  	conn := &connection{
    87  		s:      s,
    88  		conn:   wsConn,
    89  		send:   make(chan interface{}, maxPendingMessages),
    90  		fp:     NewFilterParam(),
    91  		active: 1,
    92  	}
    93  	s.addConnection(conn)
    94  }
    95  
    96  func (s *Server) Publish(parser Filterer) {
    97  	conns := s.subscribedConnections.Conns()
    98  	toNotify, msg := parser.Filter(conns)
    99  	for i, shouldNotify := range toNotify {
   100  		if !shouldNotify {
   101  			continue
   102  		}
   103  		conn := conns[i].(*connection)
   104  		if !conn.Send(msg) {
   105  			s.log.Verbo("dropping message to subscribed connection due to too many pending messages")
   106  		}
   107  	}
   108  }
   109  
   110  func (s *Server) addConnection(conn *connection) {
   111  	s.lock.Lock()
   112  	defer s.lock.Unlock()
   113  
   114  	s.conns.Add(conn)
   115  
   116  	go conn.writePump()
   117  	go conn.readPump()
   118  }
   119  
   120  func (s *Server) removeConnection(conn *connection) {
   121  	s.subscribedConnections.Remove(conn)
   122  
   123  	s.lock.Lock()
   124  	defer s.lock.Unlock()
   125  
   126  	s.conns.Remove(conn)
   127  }