stream_test.go

/*
   Copyright The containerd Authors.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package ttrpc

import (
	"context"
	"io"
	"testing"
	"time"

	"github.com/containerd/ttrpc/internal"
)

func TestStreamClient(t *testing.T) {
	var (
		ctx             = context.Background()
		server          = mustServer(t)(NewServer())
		addr, listener  = newTestListener(t)
		client, cleanup = newTestClient(t, addr)
		serviceName     = "streamService"
	)

	defer listener.Close()
	defer cleanup()

	desc := &ServiceDesc{
		Methods: map[string]Method{
			"Echo": func(_ context.Context, unmarshal func(interface{}) error) (interface{}, error) {
				var req internal.EchoPayload
				if err := unmarshal(&req); err != nil {
					return nil, err
				}
				req.Seq++
				return &req, nil
			},
		},
		Streams: map[string]Stream{
			"EchoStream": {
				Handler: func(_ context.Context, ss StreamServer) (interface{}, error) {
					for {
						var req internal.EchoPayload
						if err := ss.RecvMsg(&req); err != nil {
							if err == io.EOF {
								err = nil
							}
							return nil, err
						}
						req.Seq++
						if err := ss.SendMsg(&req); err != nil {
							return nil, err
						}
					}

				},
				StreamingClient: true,
				StreamingServer: true,
			},
		},
	}
	server.RegisterService(serviceName, desc)

	go server.Serve(ctx, listener)
	defer server.Shutdown(ctx)

	//func (c *Client) NewStream(ctx context.Context, desc *StreamDesc, service, method string) (ClientStream, error) {
	var req, resp internal.EchoPayload
	if err := client.Call(ctx, serviceName, "Echo", &req, &resp); err != nil {
		t.Fatal(err)
	}

	stream, err := client.NewStream(ctx, &StreamDesc{true, true}, serviceName, "EchoStream", nil)
	if err != nil {
		t.Fatal(err)
	}
	for i := 1; i <= 100; i++ {
		req := internal.EchoPayload{
			Seq: int64(i),
			Msg: "should be returned",
		}
		if err := stream.SendMsg(&req); err != nil {
			t.Fatalf("%d: %v", i, err)
		}
		var resp internal.EchoPayload
		if err := stream.RecvMsg(&resp); err != nil {
			t.Fatalf("%d: %v", i, err)
		}
		if resp.Seq != int64(i)+1 {
			t.Fatalf("%d: unexpected sequence value: %d, expected %d", i, resp.Seq, i+1)
		}
		if resp.Msg != req.Msg {
			t.Fatalf("%d: unexpected message: %q, expected %q", i, resp.Msg, req.Msg)
		}
	}
	if err := stream.CloseSend(); err != nil {
		t.Fatal(err)
	}

	err = stream.RecvMsg(&resp)
	if err == nil {
		t.Fatal("expected io.EOF after close send")
	}
	if err != io.EOF {
		t.Fatalf("expected io.EOF after close send, got %v", err)
	}
}

// TestStreamHandlerContextCancelOnReturn verifies that the context passed
// to a stream handler is cancelled once the handler returns. Without this,
// goroutines that the handler spawned and parked on `<-ctx.Done()` (for
// example wrappers around RecvMsg) would only wake on connection close,
// accumulating one leaked goroutine per completed stream.
func TestStreamHandlerContextCancelOnReturn(t *testing.T) {
	var (
		ctx             = context.Background()
		server          = mustServer(t)(NewServer())
		addr, listener  = newTestListener(t)
		client, cleanup = newTestClient(t, addr)
		serviceName     = "streamCancelService"
	)
	defer listener.Close()
	defer cleanup()

	handlerCtx := make(chan context.Context, 1)
	desc := &ServiceDesc{
		Streams: map[string]Stream{
			"Quick": {
				Handler: func(hctx context.Context, _ StreamServer) (interface{}, error) {
					handlerCtx <- hctx
					return nil, nil
				},
				StreamingClient: true,
				StreamingServer: true,
			},
		},
	}
	server.RegisterService(serviceName, desc)

	go server.Serve(ctx, listener)
	defer server.Shutdown(ctx)

	stream, err := client.NewStream(ctx, &StreamDesc{StreamingClient: true, StreamingServer: true}, serviceName, "Quick", nil)
	if err != nil {
		t.Fatal(err)
	}
	defer stream.CloseSend()

	var hctx context.Context
	select {
	case hctx = <-handlerCtx:
	case <-time.After(2 * time.Second):
		t.Fatal("handler did not run within 2s")
	}

	// After the handler returned, its per-stream ctx must be cancelled.
	// Pre-fix this only happened on connection close.
	select {
	case <-hctx.Done():
	case <-time.After(2 * time.Second):
		t.Fatal("handler ctx not cancelled after handler returned")
	}
}