github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/tests/integration/topic_grpc_stopper_helper_test.go (about)

     1  //go:build integration
     2  // +build integration
     3  
     4  package integration
     5  
     6  import (
     7  	"context"
     8  
     9  	"google.golang.org/grpc"
    10  
    11  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/empty"
    12  )
    13  
    14  // GrpcStopper use interceptors for stop any real grpc activity on grpc channel and return error for any calls
    15  //
    16  // Usage:
    17  //
    18  //		grpcStopper := xtest.NewGrpcStopper()
    19  //
    20  //		db, err := ydb.Open(context.Background(), connectionString,
    21  //	     ...
    22  //			ydb.With(config.WithGrpcOptions(grpc.WithChainUnaryInterceptor(grpcStopper.UnaryClientInterceptor)),
    23  //			ydb.With(config.WithGrpcOptions(grpc.WithStreamInterceptor(grpcStopper.StreamClientInterceptor)),
    24  //		),
    25  //
    26  //		grpcStopper.Stop(errors.New("test error"))
    27  //
    28  // )
    29  type GrpcStopper struct {
    30  	stopChannel empty.Chan
    31  	stopError   error
    32  }
    33  
    34  func NewGrpcStopper(closeError error) GrpcStopper {
    35  	return GrpcStopper{
    36  		stopChannel: make(empty.Chan),
    37  		stopError:   closeError,
    38  	}
    39  }
    40  
    41  func (l GrpcStopper) Stop() {
    42  	close(l.stopChannel)
    43  }
    44  
    45  func (l GrpcStopper) UnaryClientInterceptor(
    46  	ctx context.Context,
    47  	method string,
    48  	req, reply interface{},
    49  	cc *grpc.ClientConn,
    50  	invoker grpc.UnaryInvoker,
    51  	opts ...grpc.CallOption,
    52  ) error {
    53  	if isClosed(l.stopChannel) {
    54  		return l.stopError
    55  	}
    56  
    57  	resChan := make(chan error, 1)
    58  	go func() {
    59  		resChan <- invoker(ctx, method, req, reply, cc, opts...)
    60  	}()
    61  	select {
    62  	case <-l.stopChannel:
    63  		return l.stopError
    64  	case err := <-resChan:
    65  		return err
    66  	}
    67  }
    68  
    69  func (l GrpcStopper) StreamClientInterceptor(
    70  	ctx context.Context,
    71  	desc *grpc.StreamDesc,
    72  	cc *grpc.ClientConn,
    73  	method string,
    74  	streamer grpc.Streamer,
    75  	opts ...grpc.CallOption,
    76  ) (grpc.ClientStream, error) {
    77  	select {
    78  	case <-l.stopChannel:
    79  		// grpc stopped
    80  		return nil, l.stopError
    81  	default:
    82  	}
    83  
    84  	stream, err := streamer(ctx, desc, cc, method, opts...)
    85  	streamWrapper := newGrpcStopperStream(stream, l.stopChannel, l.stopError)
    86  	if stream != nil {
    87  		stream = streamWrapper
    88  	}
    89  	return stream, err
    90  }
    91  
    92  type GrpcStopperStream struct {
    93  	stopChannel empty.Chan
    94  	stopError   error
    95  	grpc.ClientStream
    96  }
    97  
    98  func newGrpcStopperStream(stream grpc.ClientStream, stopChannel empty.Chan, stopError error) GrpcStopperStream {
    99  	return GrpcStopperStream{stopChannel: stopChannel, ClientStream: stream, stopError: stopError}
   100  }
   101  
   102  func (g GrpcStopperStream) CloseSend() error {
   103  	if isClosed(g.stopChannel) {
   104  		return g.stopError
   105  	}
   106  
   107  	resChan := make(chan error, 1)
   108  	go func() {
   109  		resChan <- g.ClientStream.CloseSend()
   110  	}()
   111  	select {
   112  	case <-g.stopChannel:
   113  		return g.stopError
   114  	case err := <-resChan:
   115  		return err
   116  	}
   117  }
   118  
   119  func (g GrpcStopperStream) SendMsg(m interface{}) error {
   120  	if isClosed(g.stopChannel) {
   121  		return g.stopError
   122  	}
   123  
   124  	resChan := make(chan error, 1)
   125  	go func() {
   126  		resChan <- g.ClientStream.SendMsg(m)
   127  	}()
   128  	select {
   129  	case <-g.stopChannel:
   130  		return g.stopError
   131  	case err := <-resChan:
   132  		return err
   133  	}
   134  }
   135  
   136  func (g GrpcStopperStream) RecvMsg(m interface{}) error {
   137  	if isClosed(g.stopChannel) {
   138  		return g.stopError
   139  	}
   140  
   141  	resChan := make(chan error, 1)
   142  	go func() {
   143  		resChan <- g.ClientStream.RecvMsg(m)
   144  	}()
   145  
   146  	select {
   147  	case <-g.stopChannel:
   148  		return g.stopError
   149  	case err := <-resChan:
   150  		return err
   151  	}
   152  }
   153  
   154  func isClosed(ch empty.Chan) bool {
   155  	select {
   156  	case <-ch:
   157  		return true
   158  	default:
   159  		return false
   160  	}
   161  }