Do not return EOF directly, returning nil will send EOF to client (#142)

* Do not return EOF directly, returning nil will send EOF to client

* Fix lint

Author: temporal-nick

proxy/admin_stream_transfer.go

@@ -4,7 +4,6 @@ import (
 	"fmt"
 	"io"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"go.temporal.io/api/serviceerror"
@@ -53,7 +52,6 @@ func transferSourceToTarget(
 	targetStreamServer adminservice.AdminService_StreamWorkflowReplicationMessagesServer,
 	wg *sync.WaitGroup,
 	shutdownChan channel.ShutdownOnce,
-	sendEOFToServer *atomic.Bool,
 	directionLabel string,
 	logger log.Logger,
 ) {
@@ -78,13 +76,11 @@ func transferSourceToTarget(
 		if err == io.EOF {
 			logger.Debug("sourceStreamClient.Recv encountered EOF", tag.Error(err))
 			metrics.AdminServiceStreamTerminatedCount.WithLabelValues(directionLabel, "source").Inc()
-			sendEOFToServer.Store(true)
 			return
 		}
 
 		if err != nil {
 			logger.Error("sourceStreamClient.Recv encountered error", tag.Error(err))
-			sendEOFToServer.Store(true)
 			return
 		}
 		switch attr := resp.GetAttributes().(type) {
@@ -98,15 +94,13 @@ func transferSourceToTarget(
 					logger.Debug("targetStreamServer.Send encountered EOF", tag.Error(err))
 					metrics.AdminServiceStreamTerminatedCount.WithLabelValues(directionLabel, "target").Inc()
 				}
-				sendEOFToServer.Store(true)
 				return
 			}
 			metrics.AdminServiceStreamReqCount.WithLabelValues(directionLabel).Inc()
 		default:
 			logger.Error("sourceStreamClient.Recv encountered error", tag.Error(serviceerror.NewInternal(fmt.Sprintf(
 				"StreamWorkflowReplicationMessages encountered unknown type: %T %v", attr, attr,
 			))))
-			sendEOFToServer.Store(true)
 			return
 		}
 	}
@@ -117,7 +111,6 @@ func transferTargetToSource(
 	targetStreamServer adminservice.AdminService_StreamWorkflowReplicationMessagesServer,
 	wg *sync.WaitGroup,
 	shutdownChan channel.ShutdownOnce,
-	sendEOFToServer *atomic.Bool,
 	directionLabel string,
 	logger log.Logger,
 ) {
@@ -159,7 +152,6 @@ func transferTargetToSource(
 		if err == io.EOF {
 			logger.Debug("targetStreamServer.Recv encountered EOF", tag.Error(err))
 			metrics.AdminServiceStreamTerminatedCount.WithLabelValues(directionLabel, "target").Inc()
-			sendEOFToServer.Store(true)
 			return
 		}
 
@@ -178,15 +170,13 @@ func transferTargetToSource(
 					logger.Debug("sourceStreamClient.Send encountered EOF", tag.Error(err))
 					metrics.AdminServiceStreamTerminatedCount.WithLabelValues(directionLabel, "source").Inc()
 				}
-				sendEOFToServer.Store(true)
 				return
 			}
 			metrics.AdminServiceStreamRespCount.WithLabelValues(directionLabel).Inc()
 		default:
 			logger.Error("targetStreamServer.Recv encountered error", tag.Error(serviceerror.NewInternal(fmt.Sprintf(
 				"StreamWorkflowReplicationMessages encountered unknown type: %T %v", attr, attr,
 			))))
-			sendEOFToServer.Store(true)
 			return
 		}
 	}

proxy/adminservice.go

@@ -3,9 +3,7 @@ package proxy
 import (
 	"context"
 	"fmt"
-	"io"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"go.temporal.io/api/serviceerror"
@@ -299,18 +297,16 @@ func (s *adminServiceProxyServer) StreamWorkflowReplicationMessages(
 	// Scenario 2: Local disconnects. targetStreamServer.Recv will return EOF. This unblocks transferTargetToSource and sets shutdownChan.
 	//             transferSourceToTarget needs to be unblocked from sourceStreamClient.Recv
 	shutdownChan := channel.NewShutdownOnce()
-	sendEOFToServer := atomic.Bool{}
 	var wg sync.WaitGroup
 	wg.Add(2)
-	go transferTargetToSource(sourceStreamClient, targetStreamServer, &wg, shutdownChan, &sendEOFToServer, directionLabel, logger)
-	go transferSourceToTarget(sourceStreamClient, targetStreamServer, &wg, shutdownChan, &sendEOFToServer, directionLabel, logger)
+	go transferTargetToSource(sourceStreamClient, targetStreamServer, &wg, shutdownChan, directionLabel, logger)
+	go transferSourceToTarget(sourceStreamClient, targetStreamServer, &wg, shutdownChan, directionLabel, logger)
 	wg.Wait()
 
 	streamDuration := time.Since(streamStartTime)
 	metrics.AdminServiceStreamDuration.WithLabelValues(directionLabel).Observe(streamDuration.Seconds())
 	metrics.AdminServiceStreamsClosedCount.WithLabelValues(directionLabel).Inc()
-	if sendEOFToServer.Load() {
-		return io.EOF
-	}
+	// Do not try to transfer EOF from the source here. Just returning "nil" is sufficient to terminate the stream
+	// to the client.
 	return nil
 }

proxy/test/wiring_test.go

@@ -1,16 +1,24 @@
 package proxy
 
 import (
+	"context"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
+	"sync"
 	"testing"
 	"time"
 
+	"github.com/gogo/status"
 	"github.com/stretchr/testify/assert"
 	"go.temporal.io/api/workflowservice/v1"
 	"go.temporal.io/server/api/adminservice/v1"
 	"go.temporal.io/server/common/log"
+	"go.temporal.io/server/common/log/tag"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials/insecure"
 
 	"github.com/temporalio/s2s-proxy/config"
 )
@@ -52,6 +60,67 @@ var (
 	logger = log.NewTestLogger()
 )
 
+type hangupAdminServer struct {
+	adminservice.UnimplementedAdminServiceServer
+}
+
+func (s *hangupAdminServer) StreamWorkflowReplicationMessages(server adminservice.AdminService_StreamWorkflowReplicationMessagesServer) error {
+	go func() {
+		for {
+			_, err := server.Recv()
+			if status.Code(err) == codes.Canceled {
+				logger.Info("Client closed")
+				return
+			} else if err != nil {
+				logger.Info("Got a message with error", tag.Error(err))
+			}
+		}
+	}()
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+	go func() {
+		err := server.Send(&adminservice.StreamWorkflowReplicationMessagesResponse{})
+		if err != nil {
+			logger.Info("Got a message with error", tag.Error(err))
+		}
+		err = server.Send(&adminservice.StreamWorkflowReplicationMessagesResponse{})
+		if err != nil {
+			logger.Info("Got a message with error", tag.Error(err))
+		}
+		wg.Done()
+	}()
+	wg.Wait()
+	return nil
+}
+
+func TestEOFFromServer(t *testing.T) {
+	adminHandler := &hangupAdminServer{}
+	grpcHost := grpc.NewServer()
+	adminservice.RegisterAdminServiceServer(grpcHost, adminHandler)
+	listener, _ := net.Listen("tcp", "localhost:8566")
+	go func() {
+		_ = grpcHost.Serve(listener)
+	}()
+	time.Sleep(10 * time.Millisecond)
+	client, err := grpc.NewClient("localhost:8566", grpc.WithTransportCredentials(insecure.NewCredentials()))
+	assert.NoError(t, err)
+	adminServiceClient := adminservice.NewAdminServiceClient(client)
+	clientCtx, cancelCtx := context.WithCancel(t.Context())
+	streamServer, err := adminServiceClient.StreamWorkflowReplicationMessages(clientCtx)
+	assert.NoError(t, err)
+	err = streamServer.Send(&adminservice.StreamWorkflowReplicationMessagesRequest{})
+	assert.NoError(t, err)
+	_, err = streamServer.Recv()
+	assert.NoError(t, err)
+	_, err = streamServer.Recv()
+	assert.NoError(t, err)
+	_, err = streamServer.Recv()
+	assert.True(t, err == io.EOF, "Should have thrown io.EOF, but got %v instead! Error() returned: %v", err, err.Error())
+	grpcHost.Stop()
+	_ = listener.Close()
+	cancelCtx()
+}
+
 func TestWiringWithEchoService(t *testing.T) {
 	echoServer := newEchoServer(echoServerInfo, echoClientInfo, "EchoServer", logger, nil)
 	echoClient := newEchoServer(echoClientInfo, echoServerInfo, "EchoClient", logger, nil)