github.com/thanos-io/thanos@v0.32.5/pkg/store/storepb/inprocess.go (about)

     1  // Copyright (c) The Thanos Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  package storepb
     5  
     6  import (
     7  	"context"
     8  	"io"
     9  
    10  	"google.golang.org/grpc"
    11  )
    12  
    13  func ServerAsClient(srv StoreServer, clientReceiveBufferSize int) StoreClient {
    14  	return &serverAsClient{srv: srv, clientReceiveBufferSize: clientReceiveBufferSize}
    15  }
    16  
    17  // serverAsClient allows to use servers as clients.
    18  // NOTE: Passing CallOptions does not work - it would be needed to be implemented in grpc itself (before, after are private).
    19  type serverAsClient struct {
    20  	clientReceiveBufferSize int
    21  	srv                     StoreServer
    22  }
    23  
    24  func (s serverAsClient) Info(ctx context.Context, in *InfoRequest, _ ...grpc.CallOption) (*InfoResponse, error) {
    25  	return s.srv.Info(ctx, in)
    26  }
    27  
    28  func (s serverAsClient) LabelNames(ctx context.Context, in *LabelNamesRequest, _ ...grpc.CallOption) (*LabelNamesResponse, error) {
    29  	return s.srv.LabelNames(ctx, in)
    30  }
    31  
    32  func (s serverAsClient) LabelValues(ctx context.Context, in *LabelValuesRequest, _ ...grpc.CallOption) (*LabelValuesResponse, error) {
    33  	return s.srv.LabelValues(ctx, in)
    34  }
    35  
    36  func (s serverAsClient) Series(ctx context.Context, in *SeriesRequest, _ ...grpc.CallOption) (Store_SeriesClient, error) {
    37  	inSrv := &inProcessStream{recv: make(chan *SeriesResponse, s.clientReceiveBufferSize), err: make(chan error)}
    38  	inSrv.ctx, inSrv.cancel = context.WithCancel(ctx)
    39  	go func() {
    40  		inSrv.err <- s.srv.Series(in, inSrv)
    41  		close(inSrv.err)
    42  		close(inSrv.recv)
    43  	}()
    44  	return &inProcessClientStream{srv: inSrv}, nil
    45  }
    46  
    47  // TODO(bwplotka): Add streaming attributes, metadata etc. Currently those are disconnected. Follow up on https://github.com/grpc/grpc-go/issues/906.
    48  // TODO(bwplotka): Use this in proxy.go and receiver multi tenant proxy.
    49  type inProcessStream struct {
    50  	grpc.ServerStream
    51  
    52  	ctx    context.Context
    53  	cancel context.CancelFunc
    54  	recv   chan *SeriesResponse
    55  	err    chan error
    56  }
    57  
    58  func NewInProcessStream(ctx context.Context, bufferSize int) *inProcessStream {
    59  	return &inProcessStream{
    60  		ctx:  ctx,
    61  		recv: make(chan *SeriesResponse, bufferSize),
    62  		err:  make(chan error),
    63  	}
    64  }
    65  
    66  func (s *inProcessStream) Context() context.Context { return s.ctx }
    67  
    68  func (s *inProcessStream) Send(r *SeriesResponse) error {
    69  	select {
    70  	case <-s.ctx.Done():
    71  		return s.ctx.Err()
    72  	case s.recv <- r:
    73  		return nil
    74  	}
    75  }
    76  
    77  type inProcessClientStream struct {
    78  	grpc.ClientStream
    79  
    80  	srv *inProcessStream
    81  }
    82  
    83  func (s *inProcessClientStream) Context() context.Context { return s.srv.ctx }
    84  
    85  func (s *inProcessClientStream) CloseSend() error {
    86  	s.srv.cancel()
    87  	return nil
    88  }
    89  
    90  func (s *inProcessClientStream) Recv() (*SeriesResponse, error) {
    91  	select {
    92  	case <-s.srv.ctx.Done():
    93  		return nil, s.srv.ctx.Err()
    94  	case r, ok := <-s.srv.recv:
    95  		if !ok {
    96  			return nil, io.EOF
    97  		}
    98  		return r, nil
    99  	case err := <-s.srv.err:
   100  		if err == nil {
   101  			return nil, io.EOF
   102  		}
   103  		return nil, err
   104  	}
   105  }