Merge pull request #178 from cloudstruct/feat/blockfetch-timeout

feat: make blockfetch timeout configurable

Author: agaffney

protocol/blockfetch/blockfetch.go

@@ -1,6 +1,8 @@
 package blockfetch
 
 import (
+	"time"
+
 	"github.com/cloudstruct/go-ouroboros-network/protocol"
 )
 
@@ -67,10 +69,12 @@ type BlockFetch struct {
 }
 
 type Config struct {
-	StartBatchFunc StartBatchFunc
-	NoBlocksFunc   NoBlocksFunc
-	BlockFunc      BlockFunc
-	BatchDoneFunc  BatchDoneFunc
+	StartBatchFunc    StartBatchFunc
+	NoBlocksFunc      NoBlocksFunc
+	BlockFunc         BlockFunc
+	BatchDoneFunc     BatchDoneFunc
+	BatchStartTimeout time.Duration
+	BlockTimeout      time.Duration
 }
 
 // Callback function types
@@ -90,7 +94,10 @@ func New(protoOptions protocol.ProtocolOptions, cfg *Config) *BlockFetch {
 type BlockFetchOptionFunc func(*Config)
 
 func NewConfig(options ...BlockFetchOptionFunc) Config {
-	c := Config{}
+	c := Config{
+		BatchStartTimeout: 5 * time.Second,
+		BlockTimeout:      5 * time.Second,
+	}
 	// Apply provided options functions
 	for _, option := range options {
 		option(&c)
@@ -121,3 +128,15 @@ func WithBatchDoneFunc(BatchDoneFunc BatchDoneFunc) BlockFetchOptionFunc {
 		c.BatchDoneFunc = BatchDoneFunc
 	}
 }
+
+func WithBatchStartTimeout(timeout time.Duration) BlockFetchOptionFunc {
+	return func(c *Config) {
+		c.BatchStartTimeout = timeout
+	}
+}
+
+func WithBlockTimeout(timeout time.Duration) BlockFetchOptionFunc {
+	return func(c *Config) {
+		c.BlockTimeout = timeout
+	}
+}

protocol/blockfetch/client.go

@@ -13,9 +13,24 @@ type Client struct {
 }
 
 func NewClient(protoOptions protocol.ProtocolOptions, cfg *Config) *Client {
+	if cfg == nil {
+		tmpCfg := NewConfig()
+		cfg = &tmpCfg
+	}
 	c := &Client{
 		config: cfg,
 	}
+	// Update state map with timeouts
+	stateMap := StateMap
+	if entry, ok := stateMap[STATE_BUSY]; ok {
+		entry.Timeout = c.config.BatchStartTimeout
+		stateMap[STATE_BUSY] = entry
+	}
+	if entry, ok := stateMap[STATE_STREAMING]; ok {
+		entry.Timeout = c.config.BlockTimeout
+		stateMap[STATE_STREAMING] = entry
+	}
+	// Configure underlying Protocol
 	protoConfig := protocol.ProtocolConfig{
 		Name:                PROTOCOL_NAME,
 		ProtocolId:          PROTOCOL_ID,