trpc.group/trpc-go/trpc-go@v1.0.3/transport/server_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  	"fmt"
    19  
    20  	"trpc.group/trpc-go/trpc-go/codec"
    21  	"trpc.group/trpc-go/trpc-go/errs"
    22  	"trpc.group/trpc-go/trpc-go/internal/addrutil"
    23  )
    24  
    25  // serverStreamTransport implements ServerStreamTransport and keeps backward compatibility with the
    26  // original serverTransport.
    27  type serverStreamTransport struct {
    28  	// Keep backward compatibility with original serverTransport.
    29  	serverTransport
    30  }
    31  
    32  // NewServerStreamTransport creates a new ServerTransport, which is wrapped in serverStreamTransport
    33  // as the return ServerStreamTransport interface.
    34  func NewServerStreamTransport(opt ...ServerTransportOption) ServerStreamTransport {
    35  	s := newServerTransport(opt...)
    36  	return &serverStreamTransport{s}
    37  }
    38  
    39  // DefaultServerStreamTransport is the default ServerStreamTransport.
    40  var DefaultServerStreamTransport = NewServerStreamTransport()
    41  
    42  // ListenAndServe implements ServerTransport.
    43  // To be compatible with common RPC and stream RPC, we use serverTransport.ListenAndServe function.
    44  func (st *serverStreamTransport) ListenAndServe(ctx context.Context, opts ...ListenServeOption) error {
    45  	return st.serverTransport.ListenAndServe(ctx, opts...)
    46  }
    47  
    48  // Send is the method to send stream messages.
    49  func (st *serverStreamTransport) Send(ctx context.Context, req []byte) error {
    50  	msg := codec.Message(ctx)
    51  	raddr := msg.RemoteAddr()
    52  	laddr := msg.LocalAddr()
    53  	if raddr == nil || laddr == nil {
    54  		return errs.NewFrameError(errs.RetServerSystemErr,
    55  			fmt.Sprintf("Address is invalid, local: %s, remote: %s", laddr, raddr))
    56  	}
    57  	key := addrutil.AddrToKey(laddr, raddr)
    58  	st.serverTransport.m.RLock()
    59  	tc, ok := st.serverTransport.addrToConn[key]
    60  	st.serverTransport.m.RUnlock()
    61  	if ok && tc != nil {
    62  		if _, err := tc.rwc.Write(req); err != nil {
    63  			tc.close()
    64  			st.Close(ctx)
    65  			return err
    66  		}
    67  		return nil
    68  	}
    69  	return errs.NewFrameError(errs.RetServerSystemErr, "Can't find conn by addr")
    70  }
    71  
    72  // Close closes ServerStreamTransport, it also cleans up cached connections.
    73  func (st *serverStreamTransport) Close(ctx context.Context) {
    74  	msg := codec.Message(ctx)
    75  	key := addrutil.AddrToKey(msg.LocalAddr(), msg.RemoteAddr())
    76  	st.m.Lock()
    77  	delete(st.serverTransport.addrToConn, key)
    78  	st.m.Unlock()
    79  }