Merge pull request #182 from cloudstruct/feat/localstatequery-timeout

feat: make localstatequery timeouts configurable

Author: agaffney

protocol/localstatequery/client.go

@@ -22,13 +22,28 @@ type Client struct {
 }
 
 func NewClient(protoOptions protocol.ProtocolOptions, cfg *Config) *Client {
+	if cfg == nil {
+		tmpCfg := NewConfig()
+		cfg = &tmpCfg
+	}
 	c := &Client{
 		config:            cfg,
 		queryResultChan:   make(chan []byte),
 		acquireResultChan: make(chan error),
 		acquired:          false,
 		currentEra:        -1,
 	}
+	// Update state map with timeouts
+	stateMap := StateMap.Copy()
+	if entry, ok := stateMap[STATE_ACQUIRING]; ok {
+		entry.Timeout = c.config.AcquireTimeout
+		stateMap[STATE_ACQUIRING] = entry
+	}
+	if entry, ok := stateMap[STATE_QUERYING]; ok {
+		entry.Timeout = c.config.QueryTimeout
+		stateMap[STATE_QUERYING] = entry
+	}
+	// Configure underlying Protocol
 	protoConfig := protocol.ProtocolConfig{
 		Name:                PROTOCOL_NAME,
 		ProtocolId:          PROTOCOL_ID,
@@ -38,7 +53,7 @@ func NewClient(protoOptions protocol.ProtocolOptions, cfg *Config) *Client {
 		Role:                protocol.ProtocolRoleClient,
 		MessageHandlerFunc:  c.messageHandler,
 		MessageFromCborFunc: NewMsgFromCbor,
-		StateMap:            StateMap,
+		StateMap:            stateMap,
 		InitialState:        STATE_IDLE,
 	}
 	// Enable version-dependent features

protocol/localstatequery/localstatequery.go

@@ -1,6 +1,8 @@
 package localstatequery
 
 import (
+	"time"
+
 	"github.com/cloudstruct/go-ouroboros-network/protocol"
 )
 
@@ -89,11 +91,13 @@ type LocalStateQuery struct {
 }
 
 type Config struct {
-	AcquireFunc   AcquireFunc
-	QueryFunc     QueryFunc
-	ReleaseFunc   ReleaseFunc
-	ReAcquireFunc ReAcquireFunc
-	DoneFunc      DoneFunc
+	AcquireFunc    AcquireFunc
+	QueryFunc      QueryFunc
+	ReleaseFunc    ReleaseFunc
+	ReAcquireFunc  ReAcquireFunc
+	DoneFunc       DoneFunc
+	AcquireTimeout time.Duration
+	QueryTimeout   time.Duration
 }
 
 // Callback function types
@@ -115,7 +119,10 @@ func New(protoOptions protocol.ProtocolOptions, cfg *Config) *LocalStateQuery {
 type LocalStateQueryOptionFunc func(*Config)
 
 func NewConfig(options ...LocalStateQueryOptionFunc) Config {
-	c := Config{}
+	c := Config{
+		AcquireTimeout: 5 * time.Second,
+		QueryTimeout:   180 * time.Second,
+	}
 	// Apply provided options functions
 	for _, option := range options {
 		option(&c)
@@ -152,3 +159,15 @@ func WithDoneFunc(doneFunc DoneFunc) LocalStateQueryOptionFunc {
 		c.DoneFunc = doneFunc
 	}
 }
+
+func WithAcquireTimeout(timeout time.Duration) LocalStateQueryOptionFunc {
+	return func(c *Config) {
+		c.AcquireTimeout = timeout
+	}
+}
+
+func WithQueryTimeout(timeout time.Duration) LocalStateQueryOptionFunc {
+	return func(c *Config) {
+		c.QueryTimeout = timeout
+	}
+}