micro-client-grpc/stream.go

package grpc

import (
	"context"
	"io"
	"sync"

	"go.unistack.org/micro/v3/client"
	"go.unistack.org/micro/v3/tracer"
	"google.golang.org/grpc"
)

// Implements the streamer interface
type grpcStream struct {
	grpc.ClientStream
	context  context.Context
	err      error
	request  client.Request
	response client.Response
	close    func(err error)
	conn     *PoolConn
	sync.RWMutex
	closed bool
}

func (g *grpcStream) Context() context.Context {
	return g.context
}

func (g *grpcStream) Request() client.Request {
	return g.request
}

func (g *grpcStream) Response() client.Response {
	return g.response
}

func (g *grpcStream) Send(msg interface{}) error {
	if err := g.ClientStream.SendMsg(msg); err != nil {
		g.setError(err)
		return err
	}
	return nil
}

func (g *grpcStream) SendMsg(msg interface{}) error {
	if err := g.ClientStream.SendMsg(msg); err != nil {
		g.setError(err)
		return err
	}
	return nil
}

func (g *grpcStream) Recv(msg interface{}) (err error) {
	defer g.setError(err)

	if err = g.ClientStream.RecvMsg(msg); err != nil {
		// #202 - inconsistent gRPC stream behavior
		// the only way to tell if the stream is done is when we get a EOF on the Recv
		// here we should close the underlying gRPC ClientConn
		closeErr := g.Close()
		if err == io.EOF && closeErr != nil {
			err = closeErr
		}

		return err
	}

	return
}

func (g *grpcStream) RecvMsg(msg interface{}) (err error) {
	defer g.setError(err)

	if err = g.ClientStream.RecvMsg(msg); err != nil {
		// #202 - inconsistent gRPC stream behavior
		// the only way to tell if the stream is done is when we get a EOF on the Recv
		// here we should close the underlying gRPC ClientConn
		closeErr := g.Close()
		if err == io.EOF && closeErr != nil {
			err = closeErr
		}

		return err
	}

	return
}

func (g *grpcStream) Error() error {
	g.RLock()
	defer g.RUnlock()
	return g.err
}

func (g *grpcStream) setError(e error) {
	g.Lock()
	g.err = e
	g.Unlock()
}

// Close the gRPC send stream
// #202 - inconsistent gRPC stream behavior
// The underlying gRPC stream should not be closed here since the
// stream should still be able to receive after this function call
// TODO: should the conn be closed in another way?
func (g *grpcStream) Close() error {
	g.Lock()
	defer g.Unlock()

	if g.closed {
		return nil
	}

	if sp, ok := tracer.SpanFromContext(g.context); ok && sp != nil {
		if g.err != nil {
			sp.SetStatus(tracer.SpanStatusError, g.err.Error())
		}
		sp.Finish()
	}
	// close the connection
	g.closed = true
	g.close(g.err)
	return g.ClientStream.CloseSend()
}

func (g *grpcStream) CloseSend() error {
	g.Lock()
	defer g.Unlock()

	if g.closed {
		return nil
	}

	// close the connection
	g.closed = true
	g.close(g.err)
	return g.ClientStream.CloseSend()
}