github.com/blend/go-sdk@v1.20220411.3/grpcutil/graceful.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package grpcutil
     9  
    10  import (
    11  	"net"
    12  
    13  	"google.golang.org/grpc"
    14  
    15  	"github.com/blend/go-sdk/async"
    16  	"github.com/blend/go-sdk/logger"
    17  )
    18  
    19  // NewGraceful returns a new graceful host for a grpc server.
    20  func NewGraceful(listener net.Listener, server *grpc.Server) *Graceful {
    21  	return &Graceful{
    22  		Latch:    async.NewLatch(),
    23  		Listener: listener,
    24  		Server:   server,
    25  	}
    26  }
    27  
    28  // Graceful is a shim for graceful hosting grpc servers.
    29  type Graceful struct {
    30  	*async.Latch
    31  	Log      logger.Log
    32  	Listener net.Listener
    33  	Server   *grpc.Server
    34  }
    35  
    36  // WithLogger sets the logger.
    37  func (gz *Graceful) WithLogger(log logger.Log) *Graceful {
    38  	gz.Log = log
    39  	return gz
    40  }
    41  
    42  // Start starts the server.
    43  func (gz *Graceful) Start() error {
    44  	gz.Latch.Starting()
    45  	gz.Latch.Started()
    46  	logger.MaybeInfof(gz.Log, "grpc server starting, listening on %v %s", gz.Listener.Addr().Network(), gz.Listener.Addr().String())
    47  	return gz.Server.Serve(gz.Listener)
    48  }
    49  
    50  // Stop shuts the server down.
    51  func (gz *Graceful) Stop() error {
    52  	gz.Latch.Stopping()
    53  	logger.MaybeInfof(gz.Log, "grpc server shutting down")
    54  	gz.Server.GracefulStop()
    55  	gz.Latch.Stopped()
    56  	return nil
    57  }
    58  
    59  // IsRunning returns if the server is running.
    60  func (gz *Graceful) IsRunning() bool {
    61  	return gz.Latch.IsStarted()
    62  }
    63  
    64  // NotifyStarted returns the notify started signal.
    65  func (gz *Graceful) NotifyStarted() <-chan struct{} {
    66  	return gz.Latch.NotifyStarted()
    67  }
    68  
    69  // NotifyStopped returns the notify stopped signal.
    70  func (gz *Graceful) NotifyStopped() <-chan struct{} {
    71  	return gz.Latch.NotifyStopped()
    72  }