feat: simplify keepalive protocol interface

protocol/keepalive/client.go

@@ -115,6 +115,17 @@ func (c *Client) startTimer() {
 	c.timer = time.AfterFunc(c.config.Period, c.sendKeepAlive)
 }
 
+// Stop stops the keep-alive protocol client and cancels any pending timers.
+func (c *Client) Stop() {
+	c.timerMutex.Lock()
+	if c.timer != nil {
+		c.timer.Stop()
+		c.timer = nil
+	}
+	c.timerMutex.Unlock()
+	c.Protocol.Stop()
+}
+
 // messageHandler handles incoming protocol messages for the client.
 func (c *Client) messageHandler(msg protocol.Message) error {
 	var err error
@@ -149,9 +160,11 @@ func (c *Client) handleKeepAliveResponse(msgGeneric protocol.Message) error {
 			msg.Cookie,
 		)
 	}
-	if c.config != nil && c.config.KeepAliveResponseFunc != nil {
-		// Call the user callback function
-		return c.config.KeepAliveResponseFunc(c.callbackContext, msg.Cookie)
+
+	// Call optional notification callback if provided
+	if c.config != nil && c.config.OnKeepAliveResponseReceived != nil {
+		c.config.OnKeepAliveResponseReceived(c.callbackContext.ConnectionId, msg.Cookie)
 	}
+
 	return nil
 }

protocol/keepalive/keepalive.go

@@ -87,14 +87,13 @@ type KeepAlive struct {
 	Server *Server
 }
 
-// Config contains configuration options for the keep-alive protocol, including callbacks and timing parameters.
+// Config contains configuration options for the keep-alive protocol, including optional notification callbacks and timing parameters.
 type Config struct {
-	KeepAliveFunc         KeepAliveFunc
-	KeepAliveResponseFunc KeepAliveResponseFunc
-	DoneFunc              DoneFunc
-	Timeout               time.Duration
-	Period                time.Duration
-	Cookie                uint16
+	OnKeepAliveReceived         func(connection.ConnectionId, uint16)
+	OnKeepAliveResponseReceived func(connection.ConnectionId, uint16)
+	Timeout                     time.Duration
+	Period                      time.Duration
+	Cookie                      uint16
 }
 
 // CallbackContext provides context information to keep-alive protocol callbacks, including connection and role references.
@@ -104,15 +103,6 @@ type CallbackContext struct {
 	Server       *Server
 }
 
-// KeepAliveFunc is a callback function type for handling keep-alive messages.
-type KeepAliveFunc func(CallbackContext, uint16) error
-
-// KeepAliveResponseFunc is a callback function type for handling keep-alive response messages.
-type KeepAliveResponseFunc func(CallbackContext, uint16) error
-
-// DoneFunc is a callback function type for handling done messages.
-type DoneFunc func(CallbackContext) error
-
 // New creates and returns a new KeepAlive protocol instance using the provided protocol options and configuration.
 func New(protoOptions protocol.ProtocolOptions, cfg *Config) *KeepAlive {
 	k := &KeepAlive{
@@ -122,6 +112,26 @@ func New(protoOptions protocol.ProtocolOptions, cfg *Config) *KeepAlive {
 	return k
 }
 
+// Start starts both the client and server sides of the keep-alive protocol.
+func (k *KeepAlive) Start() {
+	if k.Client != nil {
+		k.Client.Start()
+	}
+	if k.Server != nil {
+		k.Server.Start()
+	}
+}
+
+// Stop stops both the client and server sides of the keep-alive protocol.
+func (k *KeepAlive) Stop() {
+	if k.Client != nil {
+		k.Client.Stop()
+	}
+	if k.Server != nil {
+		k.Server.Stop()
+	}
+}
+
 // KeepAliveOptionFunc is a function that modifies a Config.
 type KeepAliveOptionFunc func(*Config)
 
@@ -138,26 +148,17 @@ func NewConfig(options ...KeepAliveOptionFunc) Config {
 	return c
 }
 
-// WithKeepAliveFunc sets the KeepAliveFunc callback in the Config.
-func WithKeepAliveFunc(keepAliveFunc KeepAliveFunc) KeepAliveOptionFunc {
-	return func(c *Config) {
-		c.KeepAliveFunc = keepAliveFunc
-	}
-}
-
-// WithKeepAliveResponseFunc sets the KeepAliveResponseFunc callback in the Config.
-func WithKeepAliveResponseFunc(
-	keepAliveResponseFunc KeepAliveResponseFunc,
-) KeepAliveOptionFunc {
+// WithOnKeepAliveReceived sets an optional notification callback that is called when a keep-alive message is received (server side).
+func WithOnKeepAliveReceived(callback func(connection.ConnectionId, uint16)) KeepAliveOptionFunc {
 	return func(c *Config) {
-		c.KeepAliveResponseFunc = keepAliveResponseFunc
+		c.OnKeepAliveReceived = callback
 	}
 }
 
-// WithDoneFunc sets the DoneFunc callback in the Config.
-func WithDoneFunc(doneFunc DoneFunc) KeepAliveOptionFunc {
+// WithOnKeepAliveResponseReceived sets an optional notification callback that is called when a keep-alive response is received (client side).
+func WithOnKeepAliveResponseReceived(callback func(connection.ConnectionId, uint16)) KeepAliveOptionFunc {
 	return func(c *Config) {
-		c.DoneFunc = doneFunc
+		c.OnKeepAliveResponseReceived = callback
 	}
 }
 

protocol/keepalive/server.go

@@ -29,6 +29,10 @@ type Server struct {
 
 // NewServer creates and returns a new keep-alive protocol server with the given options and configuration.
 func NewServer(protoOptions protocol.ProtocolOptions, cfg *Config) *Server {
+	if cfg == nil {
+		tmpCfg := NewConfig()
+		cfg = &tmpCfg
+	}
 	s := &Server{
 		config: cfg,
 	}
@@ -60,7 +64,7 @@ func (s *Server) messageHandler(msg protocol.Message) error {
 	case MessageTypeKeepAlive:
 		err = s.handleKeepAlive(msg)
 	case MessageTypeDone:
-		err = s.handleDone()
+		s.handleDone()
 	default:
 		err = fmt.Errorf(
 			"%s: received unexpected message type %d",
@@ -81,28 +85,29 @@ func (s *Server) handleKeepAlive(msgGeneric protocol.Message) error {
 			"connection_id", s.callbackContext.ConnectionId.String(),
 		)
 	msg := msgGeneric.(*MsgKeepAlive)
-	if s.config != nil && s.config.KeepAliveFunc != nil {
-		// Call the user callback function
-		return s.config.KeepAliveFunc(s.callbackContext, msg.Cookie)
-	} else {
-		// Send the keep-alive response
-		resp := NewMsgKeepAliveResponse(msg.Cookie)
-		return s.SendMessage(resp)
+
+	// Call optional notification callback if provided
+	if s.config != nil && s.config.OnKeepAliveReceived != nil {
+		s.config.OnKeepAliveReceived(s.callbackContext.ConnectionId, msg.Cookie)
 	}
+
+	// Automatically send the keep-alive response
+	resp := NewMsgKeepAliveResponse(msg.Cookie)
+	return s.SendMessage(resp)
 }
 
 // handleDone processes a done message from the client and performs any necessary cleanup.
-func (s *Server) handleDone() error {
+func (s *Server) handleDone() {
 	s.Protocol.Logger().
 		Debug("done",
 			"component", "network",
 			"protocol", ProtocolName,
 			"role", "server",
 			"connection_id", s.callbackContext.ConnectionId.String(),
 		)
-	if s.config != nil && s.config.DoneFunc != nil {
-		// Call the user callback function
-		return s.config.DoneFunc(s.callbackContext)
-	}
-	return nil
+}
+
+// Stop stops the keep-alive protocol server.
+func (s *Server) Stop() {
+	s.Protocol.Stop()
 }