Return nil from Server on stream shutdown

collector/exporter/otelarrowexporter/internal/arrow/stream.go

@@ -355,12 +355,6 @@ func (s *Stream) read(_ context.Context) error {
 		// timeout.  TODO: possibly, improve to wait for no outstanding requests and then stop reading.
 		resp, err := s.client.Recv()
 		if err != nil {
-			// Once the send direction of stream is closed the server should return
-			// an error that mentions an EOF. The expected error code is codes.Unknown.
-			status, ok := status.FromError(err)
-			if ok && status.Message() == "EOF" && status.Code() == codes.Unknown {
-				return nil
-			}
 			// Note: do not wrap, contains a Status.
 			return err
 		}
@@ -369,6 +363,9 @@ func (s *Stream) read(_ context.Context) error {
 			return fmt.Errorf("process: %w", err)
 		}
 
+		if resp.BatchId == -1 && resp.StatusCode == arrowpb.StatusCode_OK {
+			return nil
+		}
 	}
 }
 
@@ -396,6 +393,12 @@ func (s *Stream) processBatchStatus(status *arrowpb.BatchStatus) error {
 	ch, ret := s.getSenderChannels(status)
 
 	if ch == nil {
+		// This indicates the server received EOF from client shutdown.
+		// This is not an error because this is an expected shutdown
+		// initiated by the client by setting max_stream_lifetime.
+		if status.BatchId == -1 && status.StatusCode == arrowpb.StatusCode_OK {
+			return nil
+		}
 		// In case getSenderChannels encounters a problem, the
 		// channel is nil.
 		return ret

collector/exporter/otelarrowexporter/internal/arrow/stream_test.go

@@ -127,7 +127,7 @@ func TestStreamGracefulShutdown(t *testing.T) {
 	maxStreamLifetime := 1 * time.Second
 	tc.stream.maxStreamLifetime = maxStreamLifetime
 
-	tc.fromTracesCall.Times(1).Return(oneBatch, nil)
+	tc.fromTracesCall.Times(2).Return(oneBatch, nil)
 	tc.closeSendCall.Times(1).Return(nil)
 
 	channel := newHealthyTestChannel()
@@ -140,11 +140,17 @@ func TestStreamGracefulShutdown(t *testing.T) {
 		defer wg.Done()
 		batch := <-channel.sent
 		channel.recv <- statusOKFor(batch.BatchId)
+
+		// mimick the server which will send a batchID
+		// of -1 after max_stream_lifetime elapses.
+		time.Sleep(maxStreamLifetime)
+		channel.recv <- statusOKFor(-1)
 	}()
 
 	err := tc.get().SendAndWait(tc.bgctx, twoTraces)
 	require.NoError(t, err)
-	// let stream get closed and send again.
+
+	// need to sleep so CloseSend will be called.
 	time.Sleep(maxStreamLifetime)
 	err = tc.get().SendAndWait(tc.bgctx, twoTraces)
 	require.Error(t, err)

collector/receiver/otelarrowreceiver/internal/arrow/arrow.go

@@ -299,6 +299,21 @@ func (r *Receiver) anyStream(serverStream anyStreamServer) (retErr error) {
 
 		if err != nil {
 			r.logStreamError(err)
+
+			// client called CloseSend()
+			if err.Error() == "EOF" {
+				status := &arrowpb.BatchStatus{
+					BatchId: -1,
+				}
+				status.StatusCode = arrowpb.StatusCode_OK
+				err = serverStream.Send(status)
+				if err != nil {
+					r.logStreamError(err)
+					return err
+				}
+				return nil
+			}
+
 			return err
 		}