Merge pull request #175 from cloudstruct/feat/keepalive-timeout

feat: make keep-alive timeout/period configurable

Author: agaffney

protocol/keepalive/client.go

@@ -13,9 +13,20 @@ 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 timeout
+	stateMap := StateMap
+	if entry, ok := stateMap[STATE_SERVER]; ok {
+		entry.Timeout = c.config.Timeout
+		stateMap[STATE_SERVER] = entry
+	}
+	// Configure underlying Protocol
 	protoConfig := protocol.ProtocolConfig{
 		Name:                PROTOCOL_NAME,
 		ProtocolId:          PROTOCOL_ID,
@@ -25,7 +36,7 @@ func NewClient(protoOptions protocol.ProtocolOptions, cfg *Config) *Client {
 		Role:                protocol.ProtocolRoleClient,
 		MessageHandlerFunc:  c.messageHandler,
 		MessageFromCborFunc: NewMsgFromCbor,
-		StateMap:            StateMap,
+		StateMap:            stateMap,
 		InitialState:        STATE_CLIENT,
 	}
 	c.Protocol = protocol.New(protoConfig)
@@ -38,7 +49,7 @@ func (c *Client) Start() {
 }
 
 func (c *Client) startTimer() {
-	c.timer = time.AfterFunc(KEEP_ALIVE_PERIOD*time.Second, func() {
+	c.timer = time.AfterFunc(c.config.Period, func() {
 		msg := NewMsgKeepAlive(0)
 		if err := c.SendMessage(msg); err != nil {
 			c.SendError(err)

protocol/keepalive/keepalive.go

@@ -1,6 +1,8 @@
 package keepalive
 
 import (
+	"time"
+
 	"github.com/cloudstruct/go-ouroboros-network/protocol"
 )
 
@@ -9,7 +11,10 @@ const (
 	PROTOCOL_ID   uint16 = 8
 
 	// Time between keep-alive probes, in seconds
-	KEEP_ALIVE_PERIOD = 60
+	DEFAULT_KEEP_ALIVE_PERIOD = 60
+
+	// Timeout for keep-alive responses, in seconds
+	DEFAULT_KEEP_ALIVE_TIMEOUT = 10
 )
 
 var (
@@ -55,6 +60,8 @@ type Config struct {
 	KeepAliveFunc         KeepAliveFunc
 	KeepAliveResponseFunc KeepAliveResponseFunc
 	DoneFunc              DoneFunc
+	Timeout               time.Duration
+	Period                time.Duration
 }
 
 // Callback function types
@@ -73,7 +80,10 @@ func New(protoOptions protocol.ProtocolOptions, cfg *Config) *KeepAlive {
 type KeepAliveOptionFunc func(*Config)
 
 func NewConfig(options ...KeepAliveOptionFunc) Config {
-	c := Config{}
+	c := Config{
+		Period:  DEFAULT_KEEP_ALIVE_PERIOD * time.Second,
+		Timeout: DEFAULT_KEEP_ALIVE_TIMEOUT * time.Second,
+	}
 	// Apply provided options functions
 	for _, option := range options {
 		option(&c)
@@ -98,3 +108,15 @@ func WithDoneFunc(doneFunc DoneFunc) KeepAliveOptionFunc {
 		c.DoneFunc = doneFunc
 	}
 }
+
+func WithTimeout(timeout time.Duration) KeepAliveOptionFunc {
+	return func(c *Config) {
+		c.Timeout = timeout
+	}
+}
+
+func WithPeriod(period time.Duration) KeepAliveOptionFunc {
+	return func(c *Config) {
+		c.Period = period
+	}
+}