feat(leiosnotify): report response delivery result (#1879)

wolf31o2 · web-flow · commit 2d3d6500f71c · 2026-07-14T12:43:58.000-07:00
Signed-off-by: Chris Gianelloni &lt;wolf31o2@blinklabs.io&gt;
diff --git a/muxer/muxer.go b/muxer/muxer.go
@@ -223,7 +223,9 @@ func (m *Muxer) RegisterProtocol(
 				if !ok {
 					return
 				}
-				if err := m.Send(msg); err != nil {
+				err := m.Send(msg)
+				msg.reportDelivery(err)
+				if err != nil {
 					m.sendError(err)
 					return
 				}
diff --git a/muxer/muxer_test.go b/muxer/muxer_test.go
@@ -27,6 +27,7 @@ import (
 	"time"
 
 	"github.com/blinklabs-io/gouroboros/muxer"
+	"github.com/stretchr/testify/require"
 	"go.uber.org/goleak"
 )
 
@@ -38,6 +39,14 @@ type mockConn struct {
 	mu       sync.Mutex
 }
 
+type failingWriteConn struct {
+	*mockConn
+}
+
+func (*failingWriteConn) Write([]byte) (int, error) {
+	return 0, errors.New("test: forced write failure")
+}
+
 func newMockConn() *mockConn {
 	return &mockConn{
 		readBuf:  &bytes.Buffer{},
@@ -408,6 +417,54 @@ func TestMuxerSendReceive(t *testing.T) {
 	}
 }
 
+func TestMuxerReportsSegmentDelivery(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
+	tests := []struct {
+		name      string
+		conn      net.Conn
+		wantError bool
+	}{
+		{
+			name: "success",
+			conn: newMockConn(),
+		},
+		{
+			name:      "write failure",
+			conn:      &failingWriteConn{mockConn: newMockConn()},
+			wantError: true,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			m := muxer.New(test.conn)
+			defer m.Stop()
+			sendChan, _, _ := m.RegisterProtocol(
+				0x01,
+				muxer.ProtocolRoleInitiator,
+			)
+			m.Start()
+
+			deliveryChan := make(chan error, 1)
+			segment := muxer.NewSegment(0x01, []byte("test"), false)
+			require.NotNil(t, segment)
+			segment.SetDeliveryChan(deliveryChan)
+			sendChan <- segment
+
+			select {
+			case err := <-deliveryChan:
+				if test.wantError {
+					require.Error(t, err)
+				} else {
+					require.NoError(t, err)
+				}
+			case <-time.After(time.Second):
+				t.Fatal("timed out waiting for segment delivery result")
+			}
+		})
+	}
+}
+
 // TestDiffusionModes tests different diffusion modes
 func TestDiffusionModes(t *testing.T) {
 	defer goleak.VerifyNone(t)
diff --git a/muxer/segment.go b/muxer/segment.go
@@ -40,7 +40,25 @@ type SegmentHeader struct {
 // the actual payload
 type Segment struct {
 	SegmentHeader
-	Payload []byte
+	Payload      []byte
+	deliveryChan chan<- error
+}
+
+// SetDeliveryChan registers a channel that receives the result of writing the
+// segment to the underlying connection. The channel must be buffered so muxer
+// shutdown cannot block on an abandoned receiver.
+func (s *Segment) SetDeliveryChan(deliveryChan chan<- error) {
+	s.deliveryChan = deliveryChan
+}
+
+func (s *Segment) reportDelivery(err error) {
+	if s.deliveryChan == nil {
+		return
+	}
+	select {
+	case s.deliveryChan <- err:
+	default:
+	}
 }
 
 // NewSegment returns a new Segment given a protocol ID, payload bytes, and whether the segment
diff --git a/protocol/leiosnotify/client_test.go b/protocol/leiosnotify/client_test.go
@@ -204,21 +204,32 @@ func TestConfig(t *testing.T) {
 
 	// Test config with options
 	requestNextCalled := false
+	responseSentCalled := false
 
 	cfg = NewConfig(
 		WithTimeout(120*time.Second),
 		WithRequestNextFunc(func(ctx CallbackContext) (protocol.Message, error) {
 			requestNextCalled = true
 			return nil, nil
 		}),
+		WithResponseSentFunc(func(
+			CallbackContext,
+			protocol.Message,
+			error,
+		) {
+			responseSentCalled = true
+		}),
 	)
 
 	assert.Equal(t, 120*time.Second, cfg.Timeout)
 	assert.NotNil(t, cfg.RequestNextFunc)
+	assert.NotNil(t, cfg.ResponseSentFunc)
 
 	// Test that callback can be invoked
 	_, _ = cfg.RequestNextFunc(CallbackContext{})
 	assert.True(t, requestNextCalled)
+	cfg.ResponseSentFunc(CallbackContext{}, nil, nil)
+	assert.True(t, responseSentCalled)
 }
 
 func TestProtocolConstants(t *testing.T) {
diff --git a/protocol/leiosnotify/leiosnotify.go b/protocol/leiosnotify/leiosnotify.go
@@ -81,6 +81,7 @@ type LeiosNotify struct {
 type Config struct {
 	NotificationFunc NotificationFunc
 	PipelineLimit    int
+	ResponseSentFunc ResponseSentFunc
 	RequestNextFunc  RequestNextFunc
 	Timeout          time.Duration
 }
@@ -100,6 +101,7 @@ type CallbackContext struct {
 // Callback function types
 type (
 	RequestNextFunc  func(CallbackContext) (protocol.Message, error)
+	ResponseSentFunc func(CallbackContext, protocol.Message, error)
 	NotificationFunc func(CallbackContext, protocol.Message) error
 )
 
@@ -189,6 +191,18 @@ func WithRequestNextFunc(
 	}
 }
 
+// WithResponseSentFunc registers a callback invoked after the server attempts
+// to send a response returned by RequestNextFunc. The callback receives the
+// send result so notification sources can commit or release delivery
+// reservations without advancing them before transport delivery.
+func WithResponseSentFunc(
+	responseSentFunc ResponseSentFunc,
+) LeiosNotifyOptionFunc {
+	return func(c *Config) {
+		c.ResponseSentFunc = responseSentFunc
+	}
+}
+
 func WithTimeout(timeout time.Duration) LeiosNotifyOptionFunc {
 	return func(c *Config) {
 		c.Timeout = timeout
diff --git a/protocol/leiosnotify/server.go b/protocol/leiosnotify/server.go
@@ -119,10 +119,11 @@ func (s *Server) handleRequestNext() error {
 			"received leios-notify NotificationRequestNext message but callback returned nil",
 		)
 	}
-	if err := s.SendMessage(resp); err != nil {
-		return err
+	sendErr := s.SendMessageAndWait(resp)
+	if s.config.ResponseSentFunc != nil {
+		s.config.ResponseSentFunc(s.callbackContext, resp, sendErr)
 	}
-	return nil
+	return sendErr
 }
 
 func (s *Server) handleDone() {
diff --git a/protocol/leiosnotify/server_test.go b/protocol/leiosnotify/server_test.go
@@ -66,6 +66,18 @@ func writeLeiosNotifyTestSegment(
 	require.NoError(t, err)
 }
 
+type leiosNotifyFailWriteConn struct {
+	net.Conn
+}
+
+var errLeiosNotifyTestWrite = errors.New(
+	"test: forced transport write failure",
+)
+
+func (leiosNotifyFailWriteConn) Write([]byte) (int, error) {
+	return 0, errLeiosNotifyTestWrite
+}
+
 func TestNewServer(t *testing.T) {
 	connId := connection.ConnectionId{
 		LocalAddr:  &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0},
@@ -112,6 +124,169 @@ func TestHandleRequestNext_CallbackIsCalled(t *testing.T) {
 	assert.True(t, called, "expected RequestNextFunc to be called")
 }
 
+func TestHandleRequestNextReportsSuccessfulSend(t *testing.T) {
+	connId := connection.ConnectionId{
+		LocalAddr:  &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0},
+		RemoteAddr: &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0},
+	}
+	connA, connB := net.Pipe()
+	defer connA.Close()
+	defer connB.Close()
+	m := muxer.New(connA)
+	defer m.Stop()
+
+	response := NewMsgBlockAnnouncement(cbor.RawMessage{0x82, 0x01, 0x02})
+	type sendResult struct {
+		msg protocol.Message
+		err error
+	}
+	resultCh := make(chan sendResult, 1)
+	cfg := NewConfig(
+		WithRequestNextFunc(func(CallbackContext) (protocol.Message, error) {
+			return response, nil
+		}),
+		WithResponseSentFunc(func(
+			_ CallbackContext,
+			msg protocol.Message,
+			err error,
+		) {
+			resultCh <- sendResult{msg: msg, err: err}
+		}),
+	)
+	server := NewServer(protocol.ProtocolOptions{
+		ConnectionId: connId,
+		Muxer:        m,
+	}, &cfg)
+	server.Start()
+	defer server.Protocol.Stop()
+	m.Start()
+
+	requestData, err := cbor.Encode(NewMsgNotificationRequestNext())
+	require.NoError(t, err)
+	writeLeiosNotifyTestSegment(
+		t,
+		connB,
+		muxer.NewSegment(ProtocolId, requestData, false),
+	)
+	require.NoError(t, connB.SetReadDeadline(time.Now().Add(time.Second)))
+	_, err = readLeiosNotifyTestSegment(t, connB)
+	require.NoError(t, err)
+
+	select {
+	case result := <-resultCh:
+		require.Same(t, response, result.msg)
+		require.NoError(t, result.err)
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for response send callback")
+	}
+}
+
+func TestHandleRequestNextReportsFailedSend(t *testing.T) {
+	connId := connection.ConnectionId{
+		LocalAddr:  &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0},
+		RemoteAddr: &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0},
+	}
+	resultCh := make(chan error, 1)
+	connA, connB := net.Pipe()
+	defer connA.Close()
+	defer connB.Close()
+	m := muxer.New(leiosNotifyFailWriteConn{Conn: connA})
+	defer m.Stop()
+	cfg := NewConfig(
+		WithRequestNextFunc(func(CallbackContext) (protocol.Message, error) {
+			return NewMsgBlockAnnouncement(cbor.RawMessage{0x81, 0x01}), nil
+		}),
+		WithResponseSentFunc(func(
+			_ CallbackContext,
+			_ protocol.Message,
+			err error,
+		) {
+			resultCh <- err
+		}),
+	)
+	server := NewServer(protocol.ProtocolOptions{
+		ConnectionId: connId,
+		Muxer:        m,
+	}, &cfg)
+	server.Start()
+	defer server.Protocol.Stop()
+	m.Start()
+
+	requestData, err := cbor.Encode(NewMsgNotificationRequestNext())
+	require.NoError(t, err)
+	writeLeiosNotifyTestSegment(
+		t,
+		connB,
+		muxer.NewSegment(ProtocolId, requestData, false),
+	)
+
+	select {
+	case sendErr := <-resultCh:
+		require.ErrorIs(t, sendErr, errLeiosNotifyTestWrite)
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for failed response send callback")
+	}
+}
+
+func TestHandleRequestNextReportsShutdownBeforeDelivery(t *testing.T) {
+	connId := connection.ConnectionId{
+		LocalAddr:  &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0},
+		RemoteAddr: &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0},
+	}
+	connA, connB := net.Pipe()
+	defer connA.Close()
+	defer connB.Close()
+	m := muxer.New(connA)
+	defer m.Stop()
+
+	requestCalled := make(chan struct{})
+	releaseRequest := make(chan struct{})
+	resultCh := make(chan error, 1)
+	cfg := NewConfig(
+		WithRequestNextFunc(func(CallbackContext) (protocol.Message, error) {
+			close(requestCalled)
+			// Keep the response from being queued until shutdown is observable.
+			<-releaseRequest
+			return NewMsgBlockAnnouncement(cbor.RawMessage{0x81, 0x01}), nil
+		}),
+		WithResponseSentFunc(func(
+			_ CallbackContext,
+			_ protocol.Message,
+			err error,
+		) {
+			resultCh <- err
+		}),
+	)
+	server := NewServer(protocol.ProtocolOptions{
+		ConnectionId: connId,
+		Muxer:        m,
+	}, &cfg)
+	server.Start()
+	m.Start()
+
+	requestData, err := cbor.Encode(NewMsgNotificationRequestNext())
+	require.NoError(t, err)
+	writeLeiosNotifyTestSegment(
+		t,
+		connB,
+		muxer.NewSegment(ProtocolId, requestData, false),
+	)
+	select {
+	case <-requestCalled:
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for request callback")
+	}
+	server.Protocol.Stop()
+	close(releaseRequest)
+
+	select {
+	case sendErr := <-resultCh:
+		require.ErrorIs(t, sendErr, protocol.ErrProtocolShuttingDown)
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for shutdown delivery callback")
+	}
+}
+
 func TestHandleRequestNext_NilCallback(t *testing.T) {
 	connId := connection.ConnectionId{
 		LocalAddr:  &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0},
diff --git a/protocol/protocol.go b/protocol/protocol.go
diff --git a/protocol/protocol_test.go b/protocol/protocol_test.go

Original file line number	Diff line number	Diff line change
`@@ -223,7 +223,9 @@ func (m *Muxer) RegisterProtocol(`
`223`	`223`	`if !ok {`
`224`	`224`	`return`
`225`	`225`	`}`
`226`		`- if err := m.Send(msg); err != nil {`
	`226`	`+ err := m.Send(msg)`
	`227`	`+ msg.reportDelivery(err)`
	`228`	`+ if err != nil {`
`227`	`229`	`m.sendError(err)`
`228`	`230`	`return`
`229`	`231`	`}`
Original file line number	Diff line number	Diff line change
`@@ -119,10 +119,11 @@ func (s *Server) handleRequestNext() error {`
`119`	`119`	`"received leios-notify NotificationRequestNext message but callback returned nil",`
`120`	`120`	`)`
`121`	`121`	`}`
`122`		`- if err := s.SendMessage(resp); err != nil {`
`123`		`- return err`
	`122`	`+ sendErr := s.SendMessageAndWait(resp)`
	`123`	`+ if s.config.ResponseSentFunc != nil {`
	`124`	`+ s.config.ResponseSentFunc(s.callbackContext, resp, sendErr)`
`124`	`125`	`}`
`125`		`- return nil`
	`126`	`+ return sendErr`
`126`	`127`	`}`
`127`	`128`
`128`	`129`	`func (s *Server) handleDone() {`