github.com/lfch/etcd-io/tests/v3@v3.0.0-20221004140520-eac99acd3e9d/functional/agent/server.go (about)

     1  // Copyright 2018 The etcd Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package agent
    16  
    17  import (
    18  	"math"
    19  	"net"
    20  	"os"
    21  	"os/exec"
    22  
    23  	"github.com/lfch/etcd-io/pkg/v3/proxy"
    24  	"github.com/lfch/etcd-io/server/v3/embed"
    25  	"github.com/lfch/etcd-io/tests/v3/functional/rpcpb"
    26  
    27  	"go.uber.org/zap"
    28  	"google.golang.org/grpc"
    29  )
    30  
    31  // Server implements "rpcpb.TransportServer"
    32  // and other etcd operations as an agent
    33  // no need to lock fields since request operations are
    34  // serialized in tester-side
    35  type Server struct {
    36  	lg *zap.Logger
    37  
    38  	grpcServer *grpc.Server
    39  
    40  	network string
    41  	address string
    42  	ln      net.Listener
    43  
    44  	rpcpb.TransportServer
    45  	last rpcpb.Operation
    46  
    47  	*rpcpb.Member
    48  	*rpcpb.Tester
    49  
    50  	etcdServer  *embed.Etcd
    51  	etcdCmd     *exec.Cmd
    52  	etcdLogFile *os.File
    53  
    54  	// forward incoming advertise URLs traffic to listen URLs
    55  	advertiseClientPortToProxy map[int]proxy.Server
    56  	advertisePeerPortToProxy   map[int]proxy.Server
    57  }
    58  
    59  // NewServer returns a new agent server.
    60  func NewServer(
    61  	lg *zap.Logger,
    62  	network string,
    63  	address string,
    64  ) *Server {
    65  	return &Server{
    66  		lg:                         lg,
    67  		network:                    network,
    68  		address:                    address,
    69  		last:                       rpcpb.Operation_NOT_STARTED,
    70  		advertiseClientPortToProxy: make(map[int]proxy.Server),
    71  		advertisePeerPortToProxy:   make(map[int]proxy.Server),
    72  	}
    73  }
    74  
    75  const (
    76  	maxRequestBytes   = 1.5 * 1024 * 1024
    77  	grpcOverheadBytes = 512 * 1024
    78  	maxStreams        = math.MaxUint32
    79  	maxSendBytes      = math.MaxInt32
    80  )
    81  
    82  // StartServe starts serving agent server.
    83  func (srv *Server) StartServe() error {
    84  	var err error
    85  	srv.ln, err = net.Listen(srv.network, srv.address)
    86  	if err != nil {
    87  		return err
    88  	}
    89  
    90  	var opts []grpc.ServerOption
    91  	opts = append(opts, grpc.MaxRecvMsgSize(int(maxRequestBytes+grpcOverheadBytes)))
    92  	opts = append(opts, grpc.MaxSendMsgSize(maxSendBytes))
    93  	opts = append(opts, grpc.MaxConcurrentStreams(maxStreams))
    94  	srv.grpcServer = grpc.NewServer(opts...)
    95  
    96  	rpcpb.RegisterTransportServer(srv.grpcServer, srv)
    97  
    98  	srv.lg.Info(
    99  		"gRPC server started",
   100  		zap.String("address", srv.address),
   101  		zap.String("listener-address", srv.ln.Addr().String()),
   102  	)
   103  	err = srv.grpcServer.Serve(srv.ln)
   104  	if err != nil {
   105  		srv.lg.Warn(
   106  			"gRPC server is stopped with error",
   107  			zap.String("address", srv.address),
   108  			zap.Error(err),
   109  		)
   110  	} else {
   111  		srv.lg.Info(
   112  			"gRPC server is stopped",
   113  			zap.String("address", srv.address),
   114  		)
   115  	}
   116  	return err
   117  }
   118  
   119  // Stop stops serving gRPC server.
   120  func (srv *Server) Stop() {
   121  	srv.lg.Info("gRPC server stopping", zap.String("address", srv.address))
   122  	srv.grpcServer.Stop()
   123  	srv.lg.Info("gRPC server stopped", zap.String("address", srv.address))
   124  }
   125  
   126  // Transport communicates with etcd tester.
   127  func (srv *Server) Transport(stream rpcpb.Transport_TransportServer) (reterr error) {
   128  	errc := make(chan error, 1)
   129  	go func() {
   130  		for {
   131  			var req *rpcpb.Request
   132  			var err error
   133  			req, err = stream.Recv()
   134  			if err != nil {
   135  				errc <- err
   136  				// TODO: handle error and retry
   137  				return
   138  			}
   139  
   140  			var resp *rpcpb.Response
   141  			resp, err = srv.handleTesterRequest(req)
   142  			if err != nil {
   143  				errc <- err
   144  				// TODO: handle error and retry
   145  				return
   146  			}
   147  
   148  			if err = stream.Send(resp); err != nil {
   149  				errc <- err
   150  				// TODO: handle error and retry
   151  				return
   152  			}
   153  		}
   154  	}()
   155  
   156  	select {
   157  	case reterr = <-errc:
   158  	case <-stream.Context().Done():
   159  		reterr = stream.Context().Err()
   160  	}
   161  	return reterr
   162  }