package agent import ( "context" "net" "sync" "time" "github.com/hashicorp/go-hclog" "golang.org/x/sync/errgroup" ) // apiServers is a wrapper around errgroup.Group for managing go routines for // long running agent components (ex: http server, dns server). If any of the // servers fail, the failed channel will be closed, which will cause the agent // to be shutdown instead of running in a degraded state. // // This struct exists as a shim for using errgroup.Group without making major // changes to Agent. In the future it may be removed and replaced with more // direct usage of errgroup.Group. type apiServers struct { logger hclog.Logger group *errgroup.Group servers []apiServer // failed channel is closed when the first server goroutines exit with a // non-nil error. failed <-chan struct{} } type apiServer struct { // Protocol supported by this server. One of: dns, http, https Protocol string // Addr the server is listening on Addr net.Addr // Run will be called in a goroutine to run the server. When any Run exits // with a non-nil error, the failed channel will be closed. Run func() error // Shutdown function used to stop the server Shutdown func(context.Context) error } // NewAPIServers returns an empty apiServers that is ready to Start servers. func NewAPIServers(logger hclog.Logger) *apiServers { group, ctx := errgroup.WithContext(context.TODO()) return &apiServers{ logger: logger, group: group, failed: ctx.Done(), } } func (s *apiServers) Start(srv apiServer) { srv.logger(s.logger).Info("Starting server") s.servers = append(s.servers, srv) s.group.Go(srv.Run) } func (s apiServer) logger(base hclog.Logger) hclog.Logger { return base.With( "protocol", s.Protocol, "address", s.Addr.String(), "network", s.Addr.Network()) } // Shutdown all the servers and log any errors as warning. Each server is given // 1 second, or until ctx is cancelled, to shutdown gracefully. func (s *apiServers) Shutdown(ctx context.Context) { shutdownGroup := new(sync.WaitGroup) for i := range s.servers { server := s.servers[i] shutdownGroup.Add(1) go func() { defer shutdownGroup.Done() logger := server.logger(s.logger) logger.Info("Stopping server") ctx, cancel := context.WithTimeout(ctx, time.Second) defer cancel() if err := server.Shutdown(ctx); err != nil { logger.Warn("Failed to stop server") } }() } s.servers = nil shutdownGroup.Wait() } // WaitForShutdown waits until all server goroutines have exited. Shutdown // must be called before WaitForShutdown, otherwise it will block forever. func (s *apiServers) WaitForShutdown() error { return s.group.Wait() }