feat: Implement Sender.CloseSend to close the request stream (#24)

examples/starwars/main.go

@@ -22,7 +22,7 @@ func append(tx s2.Sender[*s2.AppendInput], tail uint64) error {
 	if err != nil {
 		return err
 	}
-	defer rtx.Close()
+	defer rtx.CloseSend()
 
 	scanner := bufio.NewScanner(conn)
 	for scanner.Scan() {

s2/batching.go

@@ -197,6 +197,9 @@ func appendRecordBatchingWorker(
 
 	var lingerCh <-chan time.Time
 
+	// Close the input sender on exit.
+	defer sender.CloseSend() //nolint:errcheck
+
 	for {
 		if recordsToFlush.IsEmpty() {
 			lingerCh = make(<-chan time.Time) // Never
@@ -262,7 +265,7 @@ func (s *AppendRecordBatchingSender) Send(record AppendRecord) error {
 }
 
 // Close the sender gracefully.
-func (s *AppendRecordBatchingSender) Close() error {
+func (s *AppendRecordBatchingSender) CloseSend() error {
 	s.closeWorkerCancel()
 	<-s.workerExitCtx.Done()
 

s2/batching_test.go

@@ -26,6 +26,7 @@ func withMaxBatchBytes(n uint) AppendRecordBatchingConfigParam {
 
 type testAppendSessionSender struct {
 	inputs atomic.Value
+	closed atomic.Bool
 }
 
 func (s *testAppendSessionSender) Inputs(t *testing.T) []*AppendInput {
@@ -36,6 +37,10 @@ func (s *testAppendSessionSender) Inputs(t *testing.T) []*AppendInput {
 }
 
 func (s *testAppendSessionSender) Send(input *AppendInput) error {
+	if s.closed.Load() {
+		return ErrAppendRecordSenderClosed
+	}
+
 	if inputs, ok := s.inputs.Load().([]*AppendInput); ok {
 		inputs = append(inputs, input)
 		s.inputs.Store(inputs)
@@ -46,6 +51,12 @@ func (s *testAppendSessionSender) Send(input *AppendInput) error {
 	return nil
 }
 
+func (s *testAppendSessionSender) CloseSend() error {
+	s.closed.Store(true)
+
+	return nil
+}
+
 func TestAppendRecordBatchingMechanics(t *testing.T) {
 	testCases := []struct {
 		MaxBatchRecords uint
@@ -90,7 +101,7 @@ func TestAppendRecordBatchingMechanics(t *testing.T) {
 					require.NoError(t, recordSender.Send(AppendRecord{Body: []byte(body)}))
 				}
 
-				require.NoError(t, recordSender.Close())
+				require.NoError(t, recordSender.CloseSend())
 
 				i := 0
 
@@ -169,7 +180,7 @@ func TestAppendRecordBatchingLinger(t *testing.T) {
 	sendNext("large string")
 	sendNext("")
 
-	require.NoError(t, recordSender.Close())
+	require.NoError(t, recordSender.CloseSend())
 
 	expectedBatches := [][]string{
 		{"r_0", "r_1"},
@@ -207,7 +218,7 @@ func TestAppendRecordBatchingErrorSizeLimits(t *testing.T) {
 	require.NoError(t, err)
 
 	require.ErrorIs(t, recordSender.Send(record), ErrRecordTooBig)
-	require.ErrorIs(t, recordSender.Close(), ErrRecordTooBig)
+	require.ErrorIs(t, recordSender.CloseSend(), ErrRecordTooBig)
 }
 
 func TestAppendRecordBatchingAppendInputOpts(t *testing.T) {
@@ -240,7 +251,7 @@ func TestAppendRecordBatchingAppendInputOpts(t *testing.T) {
 		require.NoError(t, recordSender.Send(record))
 	}
 
-	require.NoError(t, recordSender.Close())
+	require.NoError(t, recordSender.CloseSend())
 
 	batches := batchSender.Inputs(t)
 

s2/example_test.go

@@ -388,7 +388,7 @@ func ExampleStreamClient_AppendSession() {
 	if err != nil {
 		panic(err)
 	}
-	defer recordSender.Close()
+	defer recordSender.CloseSend() //nolint:errcheck
 
 	send := make(chan error)
 
@@ -472,7 +472,7 @@ func ExampleAppendRecordBatchingSender() {
 	if err != nil {
 		panic(err)
 	}
-	defer recordSender.Close()
+	defer recordSender.CloseSend() //nolint:errcheck
 
 	records := []s2.AppendRecord{
 		{Body: []byte("my record 1")},

s2/types.go

@@ -243,6 +243,8 @@ type Receiver[T any] interface {
 type Sender[T any] interface {
 	// Block until the item has been sent.
 	Send(T) error
+	// Close the sender.
+	CloseSend() error
 }
 
 type recvInner[F, T any] struct {
@@ -275,6 +277,10 @@ func (r sendInner[F, T]) Send(f F) error {
 	return r.Client.Send(t)
 }
 
+func (r sendInner[F, T]) CloseSend() error {
+	return r.Client.CloseSend()
+}
+
 type implRetentionPolicy interface {
 	implRetentionPolicy()
 }