trpc.group/trpc-go/trpc-go@v1.0.3/transport/transport.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 is the network transport layer. It is only used for basic binary data network
    15  // communication without any business logic.
    16  // By default, there is only one pluggable ServerTransport and ClientTransport.
    17  package transport
    18  
    19  import (
    20  	"context"
    21  	"net"
    22  	"reflect"
    23  	"sync"
    24  
    25  	"trpc.group/trpc-go/trpc-go/codec"
    26  )
    27  
    28  var (
    29  	svrTrans    = make(map[string]ServerTransport)
    30  	muxSvrTrans = sync.RWMutex{}
    31  
    32  	clientTrans    = make(map[string]ClientTransport)
    33  	muxClientTrans = sync.RWMutex{}
    34  )
    35  
    36  // FramerBuilder is the alias of codec.FramerBuilder.
    37  type FramerBuilder = codec.FramerBuilder
    38  
    39  // Framer is the alias of codec.Framer.
    40  type Framer = codec.Framer
    41  
    42  // contextKey is the context key.
    43  type contextKey struct {
    44  	name string
    45  }
    46  
    47  var (
    48  	// LocalAddrContextKey is the local address context key.
    49  	LocalAddrContextKey = &contextKey{"local-addr"}
    50  
    51  	// RemoteAddrContextKey is the remote address context key.
    52  	RemoteAddrContextKey = &contextKey{"remote-addr"}
    53  )
    54  
    55  // RemoteAddrFromContext gets remote address from context.
    56  func RemoteAddrFromContext(ctx context.Context) net.Addr {
    57  	addr, ok := ctx.Value(RemoteAddrContextKey).(net.Addr)
    58  	if !ok {
    59  		return nil
    60  	}
    61  	return addr
    62  }
    63  
    64  // ServerTransport defines the server transport layer interface.
    65  type ServerTransport interface {
    66  	ListenAndServe(ctx context.Context, opts ...ListenServeOption) error
    67  }
    68  
    69  // ClientTransport defines the client transport layer interface.
    70  type ClientTransport interface {
    71  	RoundTrip(ctx context.Context, req []byte, opts ...RoundTripOption) (rsp []byte, err error)
    72  }
    73  
    74  // Handler is the process function when server transport receive a package.
    75  type Handler interface {
    76  	Handle(ctx context.Context, req []byte) (rsp []byte, err error)
    77  }
    78  
    79  // CloseHandler handles the logic after connection closed.
    80  type CloseHandler interface {
    81  	HandleClose(ctx context.Context) error
    82  }
    83  
    84  var framerBuilders = make(map[string]codec.FramerBuilder)
    85  
    86  // RegisterFramerBuilder register a codec.FramerBuilder.
    87  func RegisterFramerBuilder(name string, fb codec.FramerBuilder) {
    88  	fbv := reflect.ValueOf(fb)
    89  	if fb == nil || fbv.Kind() == reflect.Ptr && fbv.IsNil() {
    90  		panic("transport: register framerBuilders nil")
    91  	}
    92  	if name == "" {
    93  		panic("transport: register framerBuilders name empty")
    94  	}
    95  	framerBuilders[name] = fb
    96  }
    97  
    98  // RegisterServerTransport register a ServerTransport.
    99  func RegisterServerTransport(name string, t ServerTransport) {
   100  	tv := reflect.ValueOf(t)
   101  	if t == nil || tv.Kind() == reflect.Ptr && tv.IsNil() {
   102  		panic("transport: register nil server transport")
   103  	}
   104  	if name == "" {
   105  		panic("transport: register empty name of server transport")
   106  	}
   107  	muxSvrTrans.Lock()
   108  	svrTrans[name] = t
   109  	muxSvrTrans.Unlock()
   110  }
   111  
   112  // GetServerTransport gets the ServerTransport.
   113  func GetServerTransport(name string) ServerTransport {
   114  	muxSvrTrans.RLock()
   115  	t := svrTrans[name]
   116  	muxSvrTrans.RUnlock()
   117  	return t
   118  }
   119  
   120  // RegisterClientTransport register a ClientTransport.
   121  func RegisterClientTransport(name string, t ClientTransport) {
   122  	tv := reflect.ValueOf(t)
   123  	if t == nil || tv.Kind() == reflect.Ptr && tv.IsNil() {
   124  		panic("transport: register nil client transport")
   125  	}
   126  	if name == "" {
   127  		panic("transport: register empty name of client transport")
   128  	}
   129  	muxClientTrans.Lock()
   130  	clientTrans[name] = t
   131  	muxClientTrans.Unlock()
   132  }
   133  
   134  // GetClientTransport gets the ClientTransport.
   135  func GetClientTransport(name string) ClientTransport {
   136  	muxClientTrans.RLock()
   137  	t := clientTrans[name]
   138  	muxClientTrans.RUnlock()
   139  	return t
   140  }
   141  
   142  // GetFramerBuilder gets the FramerBuilder by name.
   143  func GetFramerBuilder(name string) codec.FramerBuilder {
   144  	return framerBuilders[name]
   145  }
   146  
   147  // ListenAndServe wraps and starts the default server transport.
   148  func ListenAndServe(opts ...ListenServeOption) error {
   149  	return DefaultServerTransport.ListenAndServe(context.Background(), opts...)
   150  }
   151  
   152  // RoundTrip wraps and starts the default client transport.
   153  func RoundTrip(ctx context.Context, req []byte, opts ...RoundTripOption) ([]byte, error) {
   154  	return DefaultClientTransport.RoundTrip(ctx, req, opts...)
   155  }