Avoid closing the stream if server still open

examples/grpc_client_streaming.js

@@ -97,8 +97,6 @@ export default () => {
 
   // close the client stream
   stream.end();
-
-  sleep(1);
 };
 
 const pointSender = (stream, point) => {

grpc/grpc.go

@@ -131,7 +131,7 @@ func (mi *ModuleInstance) stream(c goja.ConstructorCall) *goja.Object {
 		instanceMetrics: mi.metrics,
 		builtinMetrics:  mi.vu.State().BuiltinMetrics,
 		done:            make(chan struct{}),
-		state:           opened,
+		writingState:    opened,
 
 		writeQueueCh: make(chan message),
 

grpc/stream.go

@@ -28,7 +28,6 @@ type message struct {
 
 const (
 	opened = iota + 1
-	closing
 	closed
 )
 
@@ -51,8 +50,8 @@ type stream struct {
 
 	obj *goja.Object // the object that is given to js to interact with the stream
 
-	state int8
-	done  chan struct{}
+	writingState int8
+	done         chan struct{}
 
 	writeQueueCh chan message
 
@@ -185,17 +184,19 @@ func (s *stream) readData(wg *sync.WaitGroup) {
 			return
 		}
 
-		if len(msg) == 0 && isRegularClosing(err) {
+		if len(msg) > 0 {
+			s.queueMessage(msg)
+		}
+
+		if isRegularClosing(err) {
 			s.logger.WithError(err).Debug("stream is cancelled/finished")
 
 			s.tq.Queue(func() error {
-				return s.closeWithError(nil)
+				return s.closeWithError(err)
 			})
 
 			return
 		}
-
-		s.queueMessage(msg)
 	}
 }
 
@@ -299,7 +300,7 @@ func (s *stream) on(event string, listener func(goja.Value) (goja.Value, error))
 
 // write writes a message to the stream
 func (s *stream) write(input goja.Value) {
-	if s.state != opened {
+	if s.writingState != opened {
 		return
 	}
 
@@ -320,11 +321,13 @@ func (s *stream) write(input goja.Value) {
 
 // end closes client the stream
 func (s *stream) end() {
-	if s.state == closed || s.state == closing {
+	if s.writingState == closed {
 		return
 	}
 
-	s.state = closing
+	s.logger.Debugf("finishing stream %s writing", s.method)
+
+	s.writingState = closed
 	s.writeQueueCh <- message{isClosing: true}
 }
 
@@ -334,16 +337,23 @@ func (s *stream) closeWithError(err error) error {
 	return s.callErrorListeners(err)
 }
 
-// close changes the stream state to closed and triggers the end event listeners
+// close closes the stream and call end event listeners
+// Note: in the regular closing the io.EOF could come
 func (s *stream) close(err error) {
-	if s.state == closed {
+	if err == nil {
 		return
 	}
 
-	s.logger.WithError(err).Debug("stream is closing")
+	select {
+	case <-s.done:
+		s.logger.Debugf("stream %v is already closed", s.method)
+		return
+	default:
+	}
 
-	s.state = closed
+	s.logger.Debugf("stream %s is closing", s.method)
 	close(s.done)
+
 	s.tq.Queue(func() error {
 		return s.callEventListeners(eventEnd)
 	})
@@ -354,7 +364,7 @@ func (s *stream) close(err error) {
 }
 
 func (s *stream) callErrorListeners(e error) error {
-	if e == nil {
+	if e == nil || errors.Is(e, io.EOF) {
 		return nil
 	}
 

grpc/stream_test.go

@@ -3,6 +3,7 @@ package grpc
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/dop251/goja"
 	"github.com/grafana/xk6-grpc/grpc/testutils/grpcservice"
@@ -188,6 +189,101 @@ func TestStream_ErrorHandling(t *testing.T) {
 	)
 }
 
+// this test case is checking that everything that server sends
+// after the client finished (client.end called) is delivered to the client
+// and the end event is called
+func TestStream_ReceiveAllServerResponsesAfterEnd(t *testing.T) {
+	t.Parallel()
+
+	ts := newTestState(t)
+
+	stub := &FeatureExplorerStub{}
+
+	savedFeatures := []*grpcservice.Feature{
+		{
+			Name: "foo",
+			Location: &grpcservice.Point{
+				Latitude:  1,
+				Longitude: 2,
+			},
+		},
+		{
+			Name: "bar",
+			Location: &grpcservice.Point{
+				Latitude:  3,
+				Longitude: 4,
+			},
+		},
+	}
+
+	stub.listFeatures = func(rect *grpcservice.Rectangle, stream grpcservice.FeatureExplorer_ListFeaturesServer) error {
+		for _, feature := range savedFeatures {
+			// adding a delay to make server response "slower"
+			time.Sleep(200 * time.Millisecond)
+
+			if err := stream.Send(feature); err != nil {
+				return err
+			}
+		}
+
+		return nil
+	}
+
+	grpcservice.RegisterFeatureExplorerServer(ts.httpBin.ServerGRPC, stub)
+
+	replace := func(code string) (goja.Value, error) {
+		return ts.VU.Runtime().RunString(ts.httpBin.Replacer.Replace(code))
+	}
+
+	initString := codeBlock{
+		code: `
+		var client = new grpc.Client();
+		client.load([], "../grpc/testutils/grpcservice/route_guide.proto");`,
+	}
+	vuString := codeBlock{
+		code: `
+		client.connect("GRPCBIN_ADDR");
+		let stream = new grpc.Stream(client, "main.FeatureExplorer/ListFeatures")
+		stream.on('data', function (data) {
+			call('Feature:' + data.name);
+		});
+		stream.on('end', function () {
+			call('End called');
+		});
+
+		stream.write({
+			lo: {
+			  latitude: 1,
+			  longitude: 2,
+			},
+			hi: {
+			  latitude: 1,
+			  longitude: 2,
+			},
+		});
+		stream.end();
+		`,
+	}
+
+	val, err := replace(initString.code)
+	assertResponse(t, initString, err, val, ts)
+
+	ts.ToVUContext()
+
+	val, err = replace(vuString.code)
+
+	ts.EventLoop.WaitOnRegistered()
+
+	assertResponse(t, vuString, err, val, ts)
+
+	assert.Equal(t, ts.callRecorder.Recorded(), []string{
+		"Feature:foo",
+		"Feature:bar",
+		"End called",
+	},
+	)
+}
+
 // FeatureExplorerStub is a stub for FeatureExplorerServer
 // it has ability to override methods
 type FeatureExplorerStub struct {