trpc.group/trpc-go/trpc-go@v1.0.3/transport/transport_stream.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  package transport
    15  
    16  import (
    17  	"context"
    18  	"reflect"
    19  	"sync"
    20  )
    21  
    22  var (
    23  	clientStreamTrans    = make(map[string]ClientStreamTransport)
    24  	muxClientStreamTrans = sync.RWMutex{}
    25  
    26  	serverStreamTrans    = make(map[string]ServerStreamTransport)
    27  	muxServerStreamTrans = sync.RWMutex{}
    28  )
    29  
    30  // ClientStreamTransport is the client stream transport interface.
    31  // It's compatible with common RPC transport.
    32  type ClientStreamTransport interface {
    33  	// Send sends stream messages.
    34  	Send(ctx context.Context, req []byte, opts ...RoundTripOption) error
    35  	// Recv receives stream messages.
    36  	Recv(ctx context.Context, opts ...RoundTripOption) ([]byte, error)
    37  	// Init inits the stream.
    38  	Init(ctx context.Context, opts ...RoundTripOption) error
    39  	// Close closes stream transport, return connection to the resource pool.
    40  	Close(ctx context.Context)
    41  }
    42  
    43  // ServerStreamTransport is the server stream transport interface.
    44  // It's compatible with common RPC transport.
    45  type ServerStreamTransport interface {
    46  	// ServerTransport is used to keep compatibility with common RPC transport.
    47  	ServerTransport
    48  	// Send sends messages.
    49  	Send(ctx context.Context, req []byte) error
    50  	// Close is called when server encounters an error and cleans up.
    51  	Close(ctx context.Context)
    52  }
    53  
    54  // RegisterServerStreamTransport Registers a named ServerStreamTransport.
    55  func RegisterServerStreamTransport(name string, t ServerStreamTransport) {
    56  	tv := reflect.ValueOf(t)
    57  	if t == nil || tv.Kind() == reflect.Ptr && tv.IsNil() {
    58  		panic("transport: register nil server transport")
    59  	}
    60  	if name == "" {
    61  		panic("transport: register empty name of server transport")
    62  	}
    63  	muxServerStreamTrans.Lock()
    64  	serverStreamTrans[name] = t
    65  	muxServerStreamTrans.Unlock()
    66  
    67  }
    68  
    69  // RegisterClientStreamTransport registers a named ClientStreamTransport.
    70  func RegisterClientStreamTransport(name string, t ClientStreamTransport) {
    71  	tv := reflect.ValueOf(t)
    72  	if t == nil || tv.Kind() == reflect.Ptr && tv.IsNil() {
    73  		panic("transport: register nil client transport")
    74  	}
    75  	if name == "" {
    76  		panic("transport: register empty name of client transport")
    77  	}
    78  	muxClientStreamTrans.Lock()
    79  	clientStreamTrans[name] = t
    80  	muxClientStreamTrans.Unlock()
    81  }
    82  
    83  // GetClientStreamTransport returns ClientStreamTransport by name.
    84  func GetClientStreamTransport(name string) ClientStreamTransport {
    85  	muxClientStreamTrans.RLock()
    86  	t := clientStreamTrans[name]
    87  	muxClientStreamTrans.RUnlock()
    88  	return t
    89  }
    90  
    91  // GetServerStreamTransport returns ServerStreamTransport by name.
    92  func GetServerStreamTransport(name string) ServerStreamTransport {
    93  	muxServerStreamTrans.RLock()
    94  	t := serverStreamTrans[name]
    95  	muxServerStreamTrans.RUnlock()
    96  	return t
    97  }