algorand · algorandskiy · Jun 29, 2023 · Jun 6, 2023 · Jun 13, 2023 · Jun 13, 2023
diff --git a/catchup/fetcher_test.go b/catchup/fetcher_test.go
@@ -268,15 +268,15 @@ func (p *testUnicastPeer) Request(ctx context.Context, tag protocol.Tag, topics
 	}
 }
 
-func (p *testUnicastPeer) Respond(ctx context.Context, reqMsg network.IncomingMessage, responseTopics network.Topics) (e error) {
+func (p *testUnicastPeer) Respond(ctx context.Context, reqMsg network.IncomingMessage, outMsg network.OutgoingMessage) (e error) {
 
 	hashKey := uint64(0)
 	channel, found := p.responseChannels[hashKey]
 	if !found {
 	}
 
 	select {
-	case channel <- &network.Response{Topics: responseTopics}:
+	case channel <- &network.Response{Topics: outMsg.Topics}:
 	default:
 	}
 

diff --git a/catchup/peerSelector_test.go b/catchup/peerSelector_test.go
@@ -59,7 +59,7 @@ func (d *mockUnicastPeer) Version() string {
 func (d *mockUnicastPeer) Request(ctx context.Context, tag network.Tag, topics network.Topics) (resp *network.Response, e error) {
 	return nil, nil
 }
-func (d *mockUnicastPeer) Respond(ctx context.Context, reqMsg network.IncomingMessage, topics network.Topics) (e error) {
+func (d *mockUnicastPeer) Respond(ctx context.Context, reqMsg network.IncomingMessage, outMsg network.OutgoingMessage) (e error) {
 	return nil
 }
 

diff --git a/config/localTemplate.go b/config/localTemplate.go
@@ -517,9 +517,9 @@ type Local struct {
 	// only relevant if TxIncomingFilteringFlags is non-zero
 	TxIncomingFilterMaxSize uint64 `version[28]:"500000"`
 
-	// BlockServiceHTTPMemCap is the memory capacity in bytes which is allowed for the block service to use for HTTP block requests.
+	// BlockServiceMemCap is the memory capacity in bytes which is allowed for the block service to use for HTTP block requests.
 	// When it exceeds this capacity, it redirects the block requests to a different node
-	BlockServiceHTTPMemCap uint64 `version[28]:"500000000"`
+	BlockServiceMemCap uint64 `version[28]:"500000000"`
 }
 
 // DNSBootstrapArray returns an array of one or more DNS Bootstrap identifiers

diff --git a/config/local_defaults.go b/config/local_defaults.go
@@ -30,7 +30,7 @@ var defaultLocal = Local{
 	Archival:                                   false,
 	BaseLoggerDebugLevel:                       4,
 	BlockServiceCustomFallbackEndpoints:        "",
-	BlockServiceHTTPMemCap:                     500000000,
+	BlockServiceMemCap:                         500000000,
 	BroadcastConnectionsLimit:                  -1,
 	CadaverDirectory:                           "",
 	CadaverSizeTarget:                          0,

diff --git a/installer/config.json.example b/installer/config.json.example
@@ -9,7 +9,7 @@
     "Archival": false,
     "BaseLoggerDebugLevel": 4,
     "BlockServiceCustomFallbackEndpoints": "",
-    "BlockServiceHTTPMemCap": 500000000,
+    "BlockServiceMemCap": 500000000,
     "BroadcastConnectionsLimit": -1,
     "CadaverDirectory": "",
     "CadaverSizeTarget": 0,

diff --git a/network/wsNetwork.go b/network/wsNetwork.go
@@ -271,6 +271,10 @@ type OutgoingMessage struct {
 	Payload []byte
 	Topics  Topics
 	reason  disconnectReason // used when Action == Disconnect
+
+	// OnRelease is a function called when outgoing message, resulting from this incoming message, is released
+	// either by being sent or discarded.
+	OnRelease func()
 }
 
 // ForwardingPolicy is an enum indicating to whom we should send a message
@@ -1320,7 +1324,7 @@ func (wn *WebsocketNetwork) messageHandlerThread(peersConnectivityCheckCh <-chan
 					wn.log.Warnf("WebsocketNetwork.messageHandlerThread: WebsocketNetwork.Broadcast returned unexpected error %v", err)
 				}
 			case Respond:
-				err := msg.Sender.(*wsPeer).Respond(wn.ctx, msg, outmsg.Topics)
+				err := msg.Sender.(*wsPeer).Respond(wn.ctx, msg, outmsg)
 				if err != nil && err != wn.ctx.Err() {
 					wn.log.Warnf("WebsocketNetwork.messageHandlerThread: wsPeer.Respond returned unexpected error %v", err)
 				}

diff --git a/network/wsNetwork_test.go b/network/wsNetwork_test.go
@@ -606,7 +606,7 @@ func TestWebsocketNetworkCancel(t *testing.T) {
 	msgs[50].ctx = ctx
 
 	for _, peer := range peers {
-		peer.sendBufferHighPrio <- sendMessages{msgs}
+		peer.sendBufferHighPrio <- sendMessages{msgs: msgs}
 	}
 
 	select {
@@ -4509,3 +4509,83 @@ func TestMergePrimarySecondaryRelayAddressListsNoDedupExp(t *testing.T) {
 		assert.ElementsMatch(t, expectedRelayAddresses, mergedRelayAddresses)
 	})
 }
+
+// TestSendMessageCallbacks tests that the SendMessage callbacks are called correctly. These are currently used for
+// decrementing the number of bytes considered currently in flight for blockservice memcaps.
+func TestSendMessageCallbacks(t *testing.T) {
+	partitiontest.PartitionTest(t)
+	netA, netB, _, closeFunc := setupWebsocketNetworkAB(t, 2)
+	_ = netB
+	defer closeFunc()
+
+	var counter uint64
+	require.NotZero(t, netA.NumPeers())
+	peer := netA.peers[0]
+	// Need to create a channel so that the message doesn't get filtered out
+	netB.peers[0].makeResponseChannel(1)
+	for i := 0; i < 100; i++ {
+		randInt := crypto.RandUint64()%(128) + 1
+		atomic.AddUint64(&counter, randInt)
+		topic := MakeTopic("val", []byte("blah"))
+		callback := func() {
+			atomic.AddUint64(&counter, ^uint64(randInt-1))
+		}
+		msg := IncomingMessage{Sender: netA.peers[0], Tag: protocol.UniEnsBlockReqTag}
+		peer.Respond(context.Background(), msg, OutgoingMessage{OnRelease: callback, Topics: Topics{topic}})
+	}
+	// force it to disconnect by removing the only response channel -- this is breach of protocol.
+	netB.peers[0].getAndRemoveResponseChannel(1)
+	// confirm that we still have messages in the send buffer
+	require.NotZero(t, len(peer.sendBufferBulk))
+
+	// confirm that eventually the messages get drained during the cleanup
+	require.Eventually(t,
+		func() bool { return atomic.LoadUint64(&counter) == uint64(0) },
+		500*time.Millisecond,
+		25*time.Millisecond,
+	)
+	require.Eventually(t,
+		func() bool { return netA.NumPeers() == 0 },
+		500*time.Millisecond,
+		25*time.Millisecond,
+	)
+}
+
+func TestSendMessageCallbackDrain(t *testing.T) {
+	partitiontest.PartitionTest(t)
+
+	node := makeTestWebsocketNode(t)
+	destPeer := wsPeer{
+		closing:            make(chan struct{}),
+		sendBufferHighPrio: make(chan sendMessages, sendBufferLength),
+		sendBufferBulk:     make(chan sendMessages, sendBufferLength),
+		conn:               &nopConnSingleton,
+	}
+	node.addPeer(&destPeer)
+	node.Start()
+	defer node.Stop()
+
+	var target, counter uint64
+	// send messages to the peer that won't read them so they will sit in the sendQueue
+	for i := 0; i < 10; i++ {
+		randInt := crypto.RandUint64()%(128) + 1
+		target += randInt
+		topic := MakeTopic("val", []byte("blah"))
+		callback := func() {
+			counter += randInt
+		}
+		msg := IncomingMessage{Sender: node.peers[0], Tag: protocol.UniEnsBlockReqTag}
+		destPeer.Respond(context.Background(), msg, OutgoingMessage{OnRelease: callback, Topics: Topics{topic}})
+	}
+	require.Len(t, destPeer.sendBufferBulk, 10)
+	require.Zero(t, counter)
+	require.Positive(t, target)
+	// close the peer to trigger draining of the queue callbacks
+	destPeer.Close(time.Now().Add(time.Second))
+
+	require.Eventually(t,
+		func() bool { return target == counter },
+		2*time.Second,
+		50*time.Millisecond,
+	)
+}
diff --git a/network/wsPeer.go b/network/wsPeer.go
@@ -175,6 +175,9 @@ type Response struct {
 
 type sendMessages struct {
 	msgs []sendMessage
+
+	// onRelease function is called when the message is released either by being sent or discarded.
+	onRelease func()
 }
 
 type wsPeer struct {
@@ -312,7 +315,7 @@ type UnicastPeer interface {
 	// Version returns the matching version from network.SupportedProtocolVersions
 	Version() string
 	Request(ctx context.Context, tag Tag, topics Topics) (resp *Response, e error)
-	Respond(ctx context.Context, reqMsg IncomingMessage, topics Topics) (e error)
+	Respond(ctx context.Context, reqMsg IncomingMessage, outMsg OutgoingMessage) (e error)
 }
 
 // TCPInfoUnicastPeer exposes information about the underlying connection if available on the platform
@@ -388,15 +391,15 @@ func (wp *wsPeer) GetUnderlyingConnTCPInfo() (*util.TCPInfo, error) {
 }
 
 // Respond sends the response of a request message
-func (wp *wsPeer) Respond(ctx context.Context, reqMsg IncomingMessage, responseTopics Topics) (e error) {
+func (wp *wsPeer) Respond(ctx context.Context, reqMsg IncomingMessage, outMsg OutgoingMessage) (e error) {
 
 	// Get the hash/key of the request message
 	requestHash := hashTopics(reqMsg.Data)
 
 	// Add the request hash
 	requestHashData := make([]byte, binary.MaxVarintLen64)
 	binary.PutUvarint(requestHashData, requestHash)
-	responseTopics = append(responseTopics, Topic{key: requestHashKey, data: requestHashData})
+	responseTopics := append(outMsg.Topics, Topic{key: requestHashKey, data: requestHashData})
 
 	// Serialize the topics
 	serializedMsg := responseTopics.MarshallTopics()
@@ -411,11 +414,13 @@ func (wp *wsPeer) Respond(ctx context.Context, reqMsg IncomingMessage, responseT
 	}
 
 	select {
-	case wp.sendBufferBulk <- sendMessages{msgs: msg}:
+	case wp.sendBufferBulk <- sendMessages{msgs: msg, onRelease: outMsg.OnRelease}:
 	case <-wp.closing:
+		outMsg.OnRelease()
 		wp.net.log.Debugf("peer closing %s", wp.conn.RemoteAddr().String())
 		return
 	case <-ctx.Done():
+		outMsg.OnRelease()
 		return ctx.Err()
 	}
 	return nil
@@ -715,6 +720,9 @@ func (wp *wsPeer) handleFilterMessage(msg IncomingMessage) {
 }
 
 func (wp *wsPeer) writeLoopSend(msgs sendMessages) disconnectReason {
+	if msgs.onRelease != nil {
+		defer msgs.onRelease()
+	}
 	for _, msg := range msgs.msgs {
 		select {
 		case <-msg.ctx.Done():
@@ -923,6 +931,21 @@ func (wp *wsPeer) Close(deadline time.Time) {
 			wp.net.log.Infof("failed to CloseWithoutFlush to connection for %s", wp.conn.RemoteAddr().String())
 		}
 	}
+
+	// We need to loop through all of the messages with callbacks still in the send queue and call them
+	// to ensure that state of counters such as wsBlockBytesUsed is correct.
+L:
+	for {
+		select {
+		case msgs := <-wp.sendBufferBulk:
+			if msgs.onRelease != nil {
+				msgs.onRelease()
+			}
+		default:
+			break L
+		}
+
+	}
 	// now call all registered closers
 	for _, f := range wp.closers {
 		f()

diff --git a/rpcs/blockService.go b/rpcs/blockService.go
@@ -26,6 +26,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 
 	"github.com/gorilla/mux"
 
@@ -42,6 +43,7 @@ import (
 	"github.com/algorand/go-algorand/logging"
 	"github.com/algorand/go-algorand/network"
 	"github.com/algorand/go-algorand/protocol"
+	"github.com/algorand/go-algorand/util/metrics"
 )
 
 // BlockResponseContentType is the HTTP Content-Type header for a raw binary block
@@ -67,12 +69,21 @@ const (
 
 var errBlockServiceClosed = errors.New("block service is shutting down")
 
+const errMemoryAtCapacityPublic = "block service memory over capacity"
+
 type errMemoryAtCapacity struct{ capacity, used uint64 }
 
 func (err errMemoryAtCapacity) Error() string {
 	return fmt.Sprintf("block service memory over capacity: %d / %d", err.used, err.capacity)
 }
 
+var wsBlockMessagesDroppedCounter = metrics.MakeCounter(
+	metrics.MetricName{Name: "algod_rpcs_ws_reqs_dropped", Description: "Number of websocket block requests dropped due to memory capacity"},
+)
+var httpBlockMessagesDroppedCounter = metrics.MakeCounter(
+	metrics.MetricName{Name: "algod_rpcs_http_reqs_dropped", Description: "Number of http block requests dropped due to memory capacity"},
+)
+
 // LedgerForBlockService describes the Ledger methods used by BlockService.
 type LedgerForBlockService interface {
 	EncodedBlockCert(rnd basics.Round) (blk []byte, cert []byte, err error)
@@ -93,6 +104,7 @@ type BlockService struct {
 	closeWaitGroup          sync.WaitGroup
 	mu                      deadlock.Mutex
 	memoryUsed              uint64
+	wsMemoryUsed            uint64
 	memoryCap               uint64
 }
 
@@ -130,7 +142,7 @@ func MakeBlockService(log logging.Logger, config config.Local, ledger LedgerForB
 		fallbackEndpoints:       makeFallbackEndpoints(log, config.BlockServiceCustomFallbackEndpoints),
 		enableArchiverFallback:  config.EnableBlockServiceFallbackToArchiver,
 		log:                     log,
-		memoryCap:               config.BlockServiceHTTPMemCap,
+		memoryCap:               config.BlockServiceMemCap,
 	}
 	if service.enableService {
 		net.RegisterHTTPHandler(BlockServiceBlockPath, service)
@@ -244,6 +256,7 @@ func (bs *BlockService) ServeHTTP(response http.ResponseWriter, request *http.Re
 				response.WriteHeader(http.StatusServiceUnavailable)
 				bs.log.Debugf("ServeHTTP: returned retry-after: %v", err)
 			}
+			httpBlockMessagesDroppedCounter.Inc(nil)
 			return
 		default:
 			// unexpected error.
@@ -301,11 +314,35 @@ const datatypeUnsupportedErrMsg = "requested data type is unsupported"
 func (bs *BlockService) handleCatchupReq(ctx context.Context, reqMsg network.IncomingMessage) {
 	target := reqMsg.Sender.(network.UnicastPeer)
 	var respTopics network.Topics
+	var n uint64
 
 	defer func() {
-		target.Respond(ctx, reqMsg, respTopics)
+		outMsg := network.OutgoingMessage{Topics: respTopics}
+		if n > 0 {
+			outMsg.OnRelease = func() {
+				atomic.AddUint64(&bs.wsMemoryUsed, ^uint64(n-1))
+			}
+			atomic.AddUint64(&bs.wsMemoryUsed, (n))
+		}
+		err := target.Respond(ctx, reqMsg, outMsg)
+		if err != nil {
+			bs.log.Warnf("BlockService handleCatchupReq: failed to respond: %s", err)
+		}
 	}()
 
+	// If we are over-capacity, we will not process the request
+	// respond to sender with error message
+	memUsed := atomic.LoadUint64(&bs.wsMemoryUsed)
+	if memUsed > bs.memoryCap {
+		err := errMemoryAtCapacity{capacity: bs.memoryCap, used: memUsed}
+		bs.log.Infof("BlockService handleCatchupReq: %s", err.Error())
+		respTopics = network.Topics{
+			network.MakeTopic(network.ErrorKey, []byte(errMemoryAtCapacityPublic)),
+		}
+		wsBlockMessagesDroppedCounter.Inc(nil)
+		return
+	}
+
 	topics, err := network.UnmarshallTopics(reqMsg.Data)
 	if err != nil {
 		bs.log.Infof("BlockService handleCatchupReq: %s", err.Error())
@@ -338,7 +375,7 @@ func (bs *BlockService) handleCatchupReq(ctx context.Context, reqMsg network.Inc
 				[]byte(roundNumberParseErrMsg))}
 		return
 	}
-	respTopics = topicBlockBytes(bs.log, bs.ledger, basics.Round(round), string(requestType))
+	respTopics, n = topicBlockBytes(bs.log, bs.ledger, basics.Round(round), string(requestType))
 	return
 }
 
@@ -416,7 +453,7 @@ func (bs *BlockService) rawBlockBytes(round basics.Round) ([]byte, error) {
 	return data, err
 }
 
-func topicBlockBytes(log logging.Logger, dataLedger LedgerForBlockService, round basics.Round, requestType string) network.Topics {
+func topicBlockBytes(log logging.Logger, dataLedger LedgerForBlockService, round basics.Round, requestType string) (network.Topics, uint64) {
 	blk, cert, err := dataLedger.EncodedBlockCert(round)
 	if err != nil {
 		switch err.(type) {
@@ -425,7 +462,7 @@ func topicBlockBytes(log logging.Logger, dataLedger LedgerForBlockService, round
 			log.Infof("BlockService topicBlockBytes: %s", err)
 		}
 		return network.Topics{
-			network.MakeTopic(network.ErrorKey, []byte(blockNotAvailableErrMsg))}
+			network.MakeTopic(network.ErrorKey, []byte(blockNotAvailableErrMsg))}, 0
 	}
 	switch requestType {
 	case BlockAndCertValue:
@@ -434,10 +471,10 @@ func topicBlockBytes(log logging.Logger, dataLedger LedgerForBlockService, round
 				BlockDataKey, blk),
 			network.MakeTopic(
 				CertDataKey, cert),
-		}
+		}, uint64(len(blk) + len(cert))
 	default:
 		return network.Topics{
-			network.MakeTopic(network.ErrorKey, []byte(datatypeUnsupportedErrMsg))}
+			network.MakeTopic(network.ErrorKey, []byte(datatypeUnsupportedErrMsg))}, 0
 	}
 }