func runSignalWrapper()

in cmd/root.go [1007:1193]
164 lines of code
28 McCabe index (conditional complexity)

func runSignalWrapper(cmd *Command) (err error) {
	defer cmd.cleanup()
	ctx, cancel := context.WithCancel(cmd.Context())
	defer cancel()

	// Configure collectors before the proxy has started to ensure we are
	// collecting metrics before *ANY* AlloyDB Admin API calls are made.
	enableMetrics := !cmd.conf.DisableMetrics
	enableTraces := !cmd.conf.DisableTraces
	if cmd.conf.TelemetryProject != "" && (enableMetrics || enableTraces) {
		sd, err := stackdriver.NewExporter(stackdriver.Options{
			ProjectID:    cmd.conf.TelemetryProject,
			MetricPrefix: cmd.conf.TelemetryPrefix,
		})
		if err != nil {
			return err
		}
		if enableMetrics {
			err = sd.StartMetricsExporter()
			if err != nil {
				return err
			}
		}
		if enableTraces {
			s := trace.ProbabilitySampler(1 / float64(cmd.conf.TelemetryTracingSampleRate))
			trace.ApplyConfig(trace.Config{DefaultSampler: s})
			trace.RegisterExporter(sd)
		}
		defer func() {
			sd.Flush()
			sd.StopMetricsExporter()
		}()
	}

	shutdownCh := make(chan error)
	// watch for sigterm / sigint signals
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGTERM, syscall.SIGINT)
	go func() {
		var s os.Signal
		select {
		case s = <-signals:
		case <-ctx.Done():
			// this should only happen when the context supplied in tests in canceled
			s = syscall.SIGINT
		}
		switch s {
		case syscall.SIGINT:
			shutdownCh <- errSigInt
		case syscall.SIGTERM:
			if cmd.conf.ExitZeroOnSigterm {
				shutdownCh <- errSigTermZero
			} else {
				shutdownCh <- errSigTerm
			}
		}
	}()

	// Start the proxy asynchronously, so we can exit early if a shutdown signal is sent
	startCh := make(chan *proxy.Client)
	go func() {
		defer close(startCh)
		p, err := proxy.NewClient(ctx, cmd.dialer, cmd.logger, cmd.conf)
		if err != nil {
			shutdownCh <- fmt.Errorf("unable to start: %v", err)
			return
		}
		startCh <- p
	}()
	// Wait for either startup to finish or a signal to interupt
	var p *proxy.Client
	select {
	case err := <-shutdownCh:
		cmd.logger.Errorf("The proxy has encountered a terminal error: %v", err)
		// If running under systemd with Type=notify, it will send a message to the
		// service manager that a failure occurred and it is terminating.
		go func() {
			if _, err := daemon.SdNotify(false, daemon.SdNotifyStopping); err != nil {
				cmd.logger.Errorf("Failed to notify systemd of termination: %v", err)
			}
		}()
		return err
	case p = <-startCh:
		cmd.logger.Infof("The proxy has started successfully and is ready for new connections!")
		// If running under systemd with Type=notify, it will send a message to the
		// service manager that it is ready to handle connections now.
		go func() {
			if _, err := daemon.SdNotify(false, daemon.SdNotifyReady); err != nil {
				cmd.logger.Errorf("Failed to notify systemd of readiness: %v", err)
			}
		}()
	}
	defer func() {
		if cErr := p.Close(); cErr != nil {
			cmd.logger.Errorf("error during shutdown: %v", cErr)
			// Capture error from close to propagate it to the caller.
			err = cErr
		}
	}()

	var (
		needsHTTPServer bool
		mux             = http.NewServeMux()
		notifyStarted   = func() {}
		notifyStopped   = func() {}
	)

	if cmd.conf.Prometheus {
		needsHTTPServer = true
		e, err := prometheus.NewExporter(prometheus.Options{
			Namespace: cmd.conf.PrometheusNamespace,
		})
		if err != nil {
			return err
		}
		mux.Handle("/metrics", e)
	}

	if cmd.conf.HealthCheck {
		needsHTTPServer = true
		cmd.logger.Infof("Starting health check server at %s",
			net.JoinHostPort(cmd.conf.HTTPAddress, cmd.conf.HTTPPort))
		hc := healthcheck.NewCheck(p, cmd.logger)
		mux.HandleFunc("/startup", hc.HandleStartup)
		mux.HandleFunc("/readiness", hc.HandleReadiness)
		mux.HandleFunc("/liveness", hc.HandleLiveness)
		notifyStarted = hc.NotifyStarted
		notifyStopped = hc.NotifyStopped
	}
	defer notifyStopped()
	// Start the HTTP server if anything requiring HTTP is specified.
	if needsHTTPServer {
		go startHTTPServer(
			ctx,
			cmd.logger,
			net.JoinHostPort(cmd.conf.HTTPAddress, cmd.conf.HTTPPort),
			mux,
			shutdownCh,
		)
	}

	var (
		needsAdminServer bool
		m                = http.NewServeMux()
	)
	if cmd.conf.QuitQuitQuit {
		needsAdminServer = true
		cmd.logger.Infof("Enabling quitquitquit endpoint at localhost:%v", cmd.conf.AdminPort)
		// quitquitquit allows for shutdown on localhost only.
		var quitOnce sync.Once
		m.HandleFunc("/quitquitquit", quitquitquit(&quitOnce, shutdownCh))
	}
	if cmd.conf.Debug {
		needsAdminServer = true
		cmd.logger.Infof("Enabling pprof endpoints at localhost:%v", cmd.conf.AdminPort)
		// pprof standard endpoints
		m.HandleFunc("/debug/pprof/", pprof.Index)
		m.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
		m.HandleFunc("/debug/pprof/profile", pprof.Profile)
		m.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
		m.HandleFunc("/debug/pprof/trace", pprof.Trace)
	}
	if needsAdminServer {
		go startHTTPServer(
			ctx,
			cmd.logger,
			net.JoinHostPort("localhost", cmd.conf.AdminPort),
			m,
			shutdownCh,
		)
	}

	go func() { shutdownCh <- p.Serve(ctx, notifyStarted) }()

	err = <-shutdownCh
	switch {
	case errors.Is(err, errSigInt):
		cmd.logger.Infof("SIGINT signal received. Shutting down...")
		time.Sleep(cmd.conf.WaitBeforeClose)
	case errors.Is(err, errSigTerm):
		cmd.logger.Infof("SIGTERM signal received. Shutting down...")
		time.Sleep(cmd.conf.WaitBeforeClose)
	default:
		cmd.logger.Errorf("The proxy has encountered a terminal error: %v", err)
	}
	return err
}