github.com/blend/go-sdk@v1.20220411.3/webutil/graceful_http_server.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 webutil
     9  
    10  import (
    11  	"context"
    12  	"net"
    13  	"net/http"
    14  	"time"
    15  
    16  	"github.com/blend/go-sdk/async"
    17  	"github.com/blend/go-sdk/ex"
    18  	"github.com/blend/go-sdk/logger"
    19  )
    20  
    21  // NewGracefulHTTPServer returns a new graceful http server wrapper.
    22  func NewGracefulHTTPServer(server *http.Server, options ...GracefulHTTPServerOption) *GracefulHTTPServer {
    23  	gs := &GracefulHTTPServer{
    24  		Latch:  async.NewLatch(),
    25  		Server: server,
    26  	}
    27  	for _, option := range options {
    28  		option(gs)
    29  	}
    30  	return gs
    31  }
    32  
    33  // GracefulHTTPServerOption is an option for the graceful http server.
    34  type GracefulHTTPServerOption func(*GracefulHTTPServer)
    35  
    36  // OptGracefulHTTPServerShutdownGracePeriod sets the shutdown grace period.
    37  func OptGracefulHTTPServerShutdownGracePeriod(d time.Duration) GracefulHTTPServerOption {
    38  	return func(g *GracefulHTTPServer) { g.ShutdownGracePeriod = d }
    39  }
    40  
    41  // OptGracefulHTTPServerListener sets the server listener.
    42  func OptGracefulHTTPServerListener(listener net.Listener) GracefulHTTPServerOption {
    43  	return func(g *GracefulHTTPServer) { g.Listener = listener }
    44  }
    45  
    46  // OptGracefulHTTPServerLog sets the logger.
    47  func OptGracefulHTTPServerLog(log logger.Log) GracefulHTTPServerOption {
    48  	return func(g *GracefulHTTPServer) { g.Log = log }
    49  }
    50  
    51  // GracefulHTTPServer is a wrapper for an http server that implements the graceful interface.
    52  type GracefulHTTPServer struct {
    53  	Log                 logger.Log
    54  	Latch               *async.Latch
    55  	Server              *http.Server
    56  	ShutdownGracePeriod time.Duration
    57  	Listener            net.Listener
    58  }
    59  
    60  // Start implements graceful.Graceful.Start.
    61  // It is expected to block.
    62  func (gs *GracefulHTTPServer) Start() (err error) {
    63  	if !gs.Latch.CanStart() {
    64  		err = ex.New(async.ErrCannotStart)
    65  		return
    66  	}
    67  	gs.Latch.Starting()
    68  	gs.Latch.Started()
    69  	defer gs.Latch.Stopped()
    70  
    71  	var shutdownErr error
    72  	if gs.Listener != nil {
    73  		logger.MaybeInfof(gs.Log, "http server listening on %s", gs.Listener.Addr().String())
    74  		shutdownErr = gs.Server.Serve(gs.Listener)
    75  	} else {
    76  		logger.MaybeInfof(gs.Log, "http server listening on %s", gs.Server.Addr)
    77  		shutdownErr = gs.Server.ListenAndServe()
    78  	}
    79  	if shutdownErr != nil && shutdownErr != http.ErrServerClosed {
    80  		err = ex.New(shutdownErr)
    81  	}
    82  	return
    83  }
    84  
    85  // Stop implements graceful.Graceful.Stop.
    86  func (gs *GracefulHTTPServer) Stop() error {
    87  	if !gs.Latch.CanStop() {
    88  		return ex.New(async.ErrCannotStop)
    89  	}
    90  	gs.Latch.Stopping()
    91  	gs.Server.SetKeepAlivesEnabled(false)
    92  	ctx := context.Background()
    93  	if gs.ShutdownGracePeriod > 0 {
    94  		var cancel context.CancelFunc
    95  		ctx, cancel = context.WithTimeout(ctx, gs.ShutdownGracePeriod)
    96  		defer cancel()
    97  	}
    98  	return ex.New(gs.Server.Shutdown(ctx))
    99  }
   100  
   101  // NotifyStarted implements part of graceful.
   102  func (gs *GracefulHTTPServer) NotifyStarted() <-chan struct{} {
   103  	return gs.Latch.NotifyStarted()
   104  }
   105  
   106  // NotifyStopped implements part of graceful.
   107  func (gs *GracefulHTTPServer) NotifyStopped() <-chan struct{} {
   108  	return gs.Latch.NotifyStopped()
   109  }