perf(protocols): improve performance across all protocols

- Increase default queue sizes for better buffering
- Eliminate unnecessary goroutines for channel cleanup
- Apply optimizations consistently across all protocols

Signed-off-by: Chris Gianelloni <wolf31o2@blinklabs.io>

Author: wolf31o2

protocol/blockfetch/blockfetch.go

@@ -119,7 +119,7 @@ type Config struct {
 const MaxRecvQueueSize = 512
 
 // DefaultRecvQueueSize is the default receive queue size.
-const DefaultRecvQueueSize = 256
+const DefaultRecvQueueSize = 384
 
 // MaxPendingMessageBytes is the maximum allowed pending message bytes (5MB).
 const MaxPendingMessageBytes = 5242880

protocol/blockfetch/client.go

@@ -36,6 +36,7 @@ type Client struct {
 	blockUseCallback     bool              // Whether to use callback for blocks
 	onceStart            sync.Once         // Ensures Start is only called once
 	onceStop             sync.Once         // Ensures Stop is only called once
+	started              bool              // Whether the protocol has been started
 }
 
 // NewClient creates a new Block Fetch protocol client with the given options and configuration.
@@ -93,13 +94,8 @@ func (c *Client) Start() {
 				"protocol", ProtocolName,
 				"connection_id", c.callbackContext.ConnectionId.String(),
 			)
+		c.started = true
 		c.Protocol.Start()
-		// Start goroutine to cleanup resources on protocol shutdown
-		go func() {
-			<-c.DoneChan()
-			close(c.blockChan)
-			close(c.startBatchResultChan)
-		}()
 	})
 }
 
@@ -115,6 +111,18 @@ func (c *Client) Stop() error {
 			)
 		msg := NewMsgClientDone()
 		err = c.SendMessage(msg)
+		// Defer closing channels until protocol fully shuts down (only if started)
+		if c.started {
+			go func() {
+				<-c.DoneChan()
+				close(c.blockChan)
+				close(c.startBatchResultChan)
+			}()
+		} else {
+			// If protocol was never started, close channels immediately
+			close(c.blockChan)
+			close(c.startBatchResultChan)
+		}
 	})
 	return err
 }
@@ -222,13 +230,11 @@ func (c *Client) handleStartBatch() error {
 			"role", "client",
 			"connection_id", c.callbackContext.ConnectionId.String(),
 		)
-	// Check for shutdown
 	select {
 	case <-c.DoneChan():
 		return protocol.ErrProtocolShuttingDown
-	default:
+	case c.startBatchResultChan <- nil:
 	}
-	c.startBatchResultChan <- nil
 	return nil
 }
 
@@ -241,14 +247,12 @@ func (c *Client) handleNoBlocks() error {
 			"role", "client",
 			"connection_id", c.callbackContext.ConnectionId.String(),
 		)
-	// Check for shutdown
+	err := errors.New("block(s) not found")
 	select {
 	case <-c.DoneChan():
 		return protocol.ErrProtocolShuttingDown
-	default:
+	case c.startBatchResultChan <- err:
 	}
-	err := errors.New("block(s) not found")
-	c.startBatchResultChan <- err
 	return nil
 }
 
@@ -298,7 +302,11 @@ func (c *Client) handleBlock(msgGeneric protocol.Message) error {
 			return errors.New("received block-fetch Block message but no callback function is defined")
 		}
 	} else {
-		c.blockChan <- block
+		select {
+		case <-c.DoneChan():
+			return protocol.ErrProtocolShuttingDown
+		case c.blockChan <- block:
+		}
 	}
 	return nil
 }

protocol/blockfetch/client_test.go

@@ -207,3 +207,58 @@ func TestGetBlockNoBlocks(t *testing.T) {
 		},
 	)
 }
+
+func TestClientShutdown(t *testing.T) {
+	defer goleak.VerifyNone(t)
+	mockConn := ouroboros_mock.NewConnection(
+		ouroboros_mock.ProtocolRoleClient,
+		[]ouroboros_mock.ConversationEntry{
+			ouroboros_mock.ConversationEntryHandshakeRequestGeneric,
+			ouroboros_mock.ConversationEntryHandshakeNtNResponse,
+		},
+	)
+	asyncErrChan := make(chan error, 1)
+	go func() {
+		err := <-mockConn.(*ouroboros_mock.Connection).ErrorChan()
+		if err != nil {
+			asyncErrChan <- fmt.Errorf("received unexpected error: %w", err)
+		}
+		close(asyncErrChan)
+	}()
+	oConn, err := ouroboros.New(
+		ouroboros.WithConnection(mockConn),
+		ouroboros.WithNetworkMagic(ouroboros_mock.MockNetworkMagic),
+		ouroboros.WithNodeToNode(true),
+	)
+	if err != nil {
+		t.Fatalf("unexpected error when creating Ouroboros object: %s", err)
+	}
+	if oConn.BlockFetch() == nil {
+		t.Fatalf("BlockFetch client is nil")
+	}
+	// Start the client
+	oConn.BlockFetch().Client.Start()
+	// Stop the client
+	if err := oConn.BlockFetch().Client.Stop(); err != nil {
+		t.Fatalf("unexpected error when stopping client: %s", err)
+	}
+	// Wait for mock connection shutdown
+	select {
+	case err, ok := <-asyncErrChan:
+		if ok {
+			t.Fatal(err.Error())
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatalf("did not complete within timeout")
+	}
+	// Close Ouroboros connection
+	if err := oConn.Close(); err != nil {
+		t.Fatalf("unexpected error when closing Ouroboros object: %s", err)
+	}
+	// Wait for connection shutdown
+	select {
+	case <-oConn.ErrorChan():
+	case <-time.After(10 * time.Second):
+		t.Errorf("did not shutdown within timeout")
+	}
+}

protocol/chainsync/chainsync.go

@@ -223,8 +223,8 @@ type Config struct {
 const (
 	MaxPipelineLimit       = 100    // Max pipelined requests
 	MaxRecvQueueSize       = 100    // Max receive queue size (messages)
-	DefaultPipelineLimit   = 50     // Default pipeline limit
-	DefaultRecvQueueSize   = 50     // Default queue size
+	DefaultPipelineLimit   = 75     // Default pipeline limit
+	DefaultRecvQueueSize   = 75     // Default queue size
 	MaxPendingMessageBytes = 102400 // Max pending message bytes (100KB)
 )
 

protocol/chainsync/client.go

@@ -34,10 +34,8 @@ type Client struct {
 	readyForNextBlockChan    chan bool
 	onceStart                sync.Once
 	onceStop                 sync.Once
+	started                  bool
 	syncPipelinedRequestNext int
-
-	// waitingForCurrentTipChan will process all the requests for the current tip until the channel
-	// is empty.
 	//
 	// want* only processes one request per message reply received from the server. If the message
 	// request fails, it is the responsibility of the caller to clear the channel.
@@ -126,12 +124,8 @@ func (c *Client) Start() {
 				"protocol", ProtocolName,
 				"connection_id", c.callbackContext.ConnectionId.String(),
 			)
+		c.started = true
 		c.Protocol.Start()
-		// Start goroutine to cleanup resources on protocol shutdown
-		go func() {
-			<-c.DoneChan()
-			close(c.readyForNextBlockChan)
-		}()
 	})
 }
 
@@ -151,6 +145,16 @@ func (c *Client) Stop() error {
 		if err = c.SendMessage(msg); err != nil {
 			return
 		}
+		// Defer closing channel until protocol fully shuts down (only if started)
+		if c.started {
+			go func() {
+				<-c.DoneChan()
+				close(c.readyForNextBlockChan)
+			}()
+		} else {
+			// If protocol was never started, close channel immediately
+			close(c.readyForNextBlockChan)
+		}
 	})
 	return err
 }
@@ -721,14 +725,22 @@ func (c *Client) handleRollForward(msgGeneric protocol.Message) error {
 	if callbackErr != nil {
 		if errors.Is(callbackErr, ErrStopSyncProcess) {
 			// Signal that we're cancelling the sync
-			c.readyForNextBlockChan <- false
+			select {
+			case <-c.DoneChan():
+				return protocol.ErrProtocolShuttingDown
+			case c.readyForNextBlockChan <- false:
+			}
 			return nil
 		} else {
 			return callbackErr
 		}
 	}
 	// Signal that we're ready for the next block
-	c.readyForNextBlockChan <- true
+	select {
+	case <-c.DoneChan():
+		return protocol.ErrProtocolShuttingDown
+	case c.readyForNextBlockChan <- true:
+	}
 	return nil
 }
 
@@ -752,15 +764,23 @@ func (c *Client) handleRollBackward(msgGeneric protocol.Message) error {
 		if callbackErr := c.config.RollBackwardFunc(c.callbackContext, msgRollBackward.Point, msgRollBackward.Tip); callbackErr != nil {
 			if errors.Is(callbackErr, ErrStopSyncProcess) {
 				// Signal that we're cancelling the sync
-				c.readyForNextBlockChan <- false
+				select {
+				case <-c.DoneChan():
+					return protocol.ErrProtocolShuttingDown
+				case c.readyForNextBlockChan <- false:
+				}
 				return nil
 			} else {
 				return callbackErr
 			}
 		}
 	}
 	// Signal that we're ready for the next block
-	c.readyForNextBlockChan <- true
+	select {
+	case <-c.DoneChan():
+		return protocol.ErrProtocolShuttingDown
+	case c.readyForNextBlockChan <- true:
+	}
 	return nil
 }
 

protocol/chainsync/client_test.go

@@ -80,6 +80,10 @@ func runTest(
 	}()
 	// Run test inner function
 	innerFunc(t, oConn)
+	// Stop the client to clean up goroutines
+	if client := oConn.ChainSync().Client; client != nil {
+		client.Stop()
+	}
 	// Wait for mock connection shutdown
 	select {
 	case err, ok := <-asyncErrChan:
@@ -275,3 +279,58 @@ func TestGetAvailableBlockRange(t *testing.T) {
 		},
 	)
 }
+
+func TestClientShutdown(t *testing.T) {
+	defer goleak.VerifyNone(t)
+	mockConn := ouroboros_mock.NewConnection(
+		ouroboros_mock.ProtocolRoleClient,
+		[]ouroboros_mock.ConversationEntry{
+			ouroboros_mock.ConversationEntryHandshakeRequestGeneric,
+			ouroboros_mock.ConversationEntryHandshakeNtCResponse,
+		},
+	)
+	asyncErrChan := make(chan error, 1)
+	go func() {
+		err := <-mockConn.(*ouroboros_mock.Connection).ErrorChan()
+		if err != nil {
+			asyncErrChan <- fmt.Errorf("received unexpected error: %w", err)
+		}
+		close(asyncErrChan)
+	}()
+	oConn, err := ouroboros.New(
+		ouroboros.WithConnection(mockConn),
+		ouroboros.WithNetworkMagic(ouroboros_mock.MockNetworkMagic),
+		ouroboros.WithChainSyncConfig(chainsync.NewConfig()),
+	)
+	if err != nil {
+		t.Fatalf("unexpected error when creating Ouroboros object: %s", err)
+	}
+	if oConn.ChainSync() == nil {
+		t.Fatalf("ChainSync client is nil")
+	}
+	// Start the client
+	oConn.ChainSync().Client.Start()
+	// Stop the client
+	if err := oConn.ChainSync().Client.Stop(); err != nil {
+		t.Fatalf("unexpected error when stopping client: %s", err)
+	}
+	// Wait for mock connection shutdown
+	select {
+	case err, ok := <-asyncErrChan:
+		if ok {
+			t.Fatal(err.Error())
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatalf("did not complete within timeout")
+	}
+	// Close Ouroboros connection
+	if err := oConn.Close(); err != nil {
+		t.Fatalf("unexpected error when closing Ouroboros object: %s", err)
+	}
+	// Wait for connection shutdown
+	select {
+	case <-oConn.ErrorChan():
+	case <-time.After(10 * time.Second):
+		t.Errorf("did not shutdown within timeout")
+	}
+}

protocol/keepalive/client.go

@@ -30,6 +30,7 @@ type Client struct {
 	timer           *time.Timer
 	timerMutex      sync.Mutex
 	onceStart       sync.Once
+	onceStop        sync.Once
 }
 
 // NewClient creates and returns a new keep-alive protocol client with the given options and configuration.
@@ -93,6 +94,22 @@ func (c *Client) Start() {
 	})
 }
 
+// Stop stops the KeepAlive client protocol
+func (c *Client) Stop() error {
+	var err error
+	c.onceStop.Do(func() {
+		c.Protocol.Logger().
+			Debug("stopping client protocol",
+				"component", "network",
+				"protocol", ProtocolName,
+				"connection_id", c.callbackContext.ConnectionId.String(),
+			)
+		msg := NewMsgDone()
+		err = c.SendMessage(msg)
+	})
+	return err
+}
+
 // sendKeepAlive sends a keep-alive message and schedules the next one.
 func (c *Client) sendKeepAlive() {
 	msg := NewMsgKeepAlive(c.config.Cookie)