github.com/gofunct/common@v0.0.0-20190131174352-fd058c7fbf22/pkg/transport/engine/engine.go (about)

     1  package engine
     2  
     3  import (
     4  	"context"
     5  	"github.com/gofunct/common/pkg/transport/api"
     6  	"github.com/gofunct/common/pkg/transport/config"
     7  	"github.com/pkg/errors"
     8  	"github.com/soheilhy/cmux"
     9  	"golang.org/x/sync/errgroup"
    10  	"google.golang.org/grpc/grpclog"
    11  	"net"
    12  	"os"
    13  	"os/signal"
    14  	"reflect"
    15  	"syscall"
    16  )
    17  
    18  // Engine is the framework instance.
    19  type Engine struct {
    20  	*config.Config
    21  	cancelFunc func()
    22  }
    23  
    24  // New creates a server intstance.
    25  func New(opts ...Option) *Engine {
    26  	return &Engine{
    27  		Config: createConfig(opts),
    28  	}
    29  }
    30  
    31  // Serve starts gRPC and Gateway servers.
    32  func (e *Engine) Serve() error {
    33  	var (
    34  		grpcServer, gatewayServer, muxServer api.Interface
    35  		grpcLis, gatewayLis, internalLis     net.Listener
    36  		err                                  error
    37  	)
    38  
    39  	if e.GrpcAddr != nil && e.GatewayAddr != nil && reflect.DeepEqual(e.GrpcAddr, e.GatewayAddr) {
    40  		lis, err := e.GrpcAddr.CreateListener()
    41  		if err != nil {
    42  			return errors.Wrap(err, "failed to listen network for servers")
    43  		}
    44  		mux := cmux.New(lis)
    45  		muxServer = NewMuxServer(mux, lis)
    46  		grpcLis = mux.Match(cmux.HTTP2HeaderField("content-type", "application/grpc"))
    47  		gatewayLis = mux.Match(cmux.HTTP2(), cmux.HTTP1Fast())
    48  	}
    49  
    50  	// Setup servers
    51  	grpcServer = NewGrpcServer(e.Config)
    52  
    53  	// Setup listeners
    54  	if grpcLis == nil && e.GrpcAddr != nil {
    55  		grpcLis, err = e.GrpcAddr.CreateListener()
    56  		if err != nil {
    57  			return errors.Wrap(err, "failed to listen network for gRPC server")
    58  		}
    59  		defer grpcLis.Close()
    60  	}
    61  
    62  	if e.GatewayAddr != nil {
    63  		gatewayServer = NewGatewayServer(e.Config)
    64  		internalLis, err = e.GrpcInternalAddr.CreateListener()
    65  		if err != nil {
    66  			return errors.Wrap(err, "failed to listen network for gRPC server internal")
    67  		}
    68  		defer internalLis.Close()
    69  	}
    70  
    71  	if gatewayLis == nil && e.GatewayAddr != nil {
    72  		gatewayLis, err = e.GatewayAddr.CreateListener()
    73  		if err != nil {
    74  			return errors.Wrap(err, "failed to listen network for gateway server")
    75  		}
    76  		defer gatewayLis.Close()
    77  	}
    78  
    79  	// Start servers
    80  	eg, ctx := errgroup.WithContext(context.Background())
    81  	ctx, e.cancelFunc = context.WithCancel(ctx)
    82  
    83  	if internalLis != nil {
    84  		eg.Go(func() error { return grpcServer.Serve(internalLis) })
    85  	}
    86  	if grpcLis != nil {
    87  		eg.Go(func() error { return grpcServer.Serve(grpcLis) })
    88  	}
    89  	if gatewayLis != nil {
    90  		eg.Go(func() error { return gatewayServer.Serve(gatewayLis) })
    91  	}
    92  	if muxServer != nil {
    93  		eg.Go(func() error { return muxServer.Serve(nil) })
    94  	}
    95  
    96  	eg.Go(func() error { return e.watchShutdownSignal(ctx) })
    97  
    98  	select {
    99  	case <-ctx.Done():
   100  		for _, s := range []api.Interface{gatewayServer, grpcServer, muxServer} {
   101  			if s != nil {
   102  				s.Shutdown()
   103  			}
   104  		}
   105  	}
   106  
   107  	err = eg.Wait()
   108  
   109  	return errors.WithStack(err)
   110  }
   111  
   112  // Shutdown closes servers.
   113  func (e *Engine) Shutdown() {
   114  	if e.cancelFunc != nil {
   115  		e.cancelFunc()
   116  	} else {
   117  		grpclog.Warning("the server has been started yet")
   118  	}
   119  }
   120  
   121  func (e *Engine) watchShutdownSignal(ctx context.Context) error {
   122  	sdCh := make(chan os.Signal, 1)
   123  	defer close(sdCh)
   124  	defer signal.Stop(sdCh)
   125  	signal.Notify(sdCh, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
   126  	select {
   127  	case <-sdCh:
   128  		e.Shutdown()
   129  	case <-ctx.Done():
   130  		// no-op
   131  	}
   132  	return nil
   133  }