Graceful shutdown in Go, pgx as db driver, chi as router, zerolog as logger, context to handle cancellation, envconfig to validates env value

main.go

package main

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/go-chi/chi/v5"
	"github.com/go-chi/chi/v5/middleware"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/joho/godotenv"
	"github.com/kelseyhightower/envconfig"
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"
)

type Env struct {
	DatabaseURL string `envconfig:"DATABASE_URL" required:"true"`
	Port        string `envconfig:"PORT" default:"8080"`
}

func main() {
	log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr, TimeFormat: time.RFC3339})

	err := godotenv.Load()
	if err != nil {
		log.Fatal().Err(err).Msg("error loading .env file")
	}

	var env Env
	err = envconfig.Process("", &env)
	// Context that is cancelled on SIGINT or SIGTERM.
	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()

	// Connect to the database.
	db, err := pgxpool.New(ctx, env.DatabaseURL)
	if err != nil {
		log.Error().Err(err).Msg("Unable to create connection pool")
		return
	}
	defer db.Close()
	log.Info().Msg("Connected to database")

	r := chi.NewRouter()
	r.Use(middleware.Logger)

	// Routes.
	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
		var email string
		var password string
		err = db.QueryRow(r.Context(), "select email, password from users").Scan(&email, &password)
		if err != nil {
			log.Fatal().Err(err).Msg("queryRow failed")
		}

		w.Header().Set("Content-Type", "application/json")
		_, err = w.Write([]byte(`{"email":"` + email + `","password":"` + password + `"}`))
		if err != nil {
			log.Error().Err(err).Msg("failed to write response")
		}
	})

	// create server
	server := &http.Server{
		Addr:         fmt.Sprintf(":%s", env.Port),
		Handler:      r,
		ReadTimeout:  10 * time.Second,
		WriteTimeout: 30 * time.Second,
		IdleTimeout:  60 * time.Second,
	}

	// Start the server in a background goroutine.
	// Send any unexpected error to the channel.
	serverErrors := make(chan error, 1)
	go func() {
		log.Info().Msgf("Server listening on %s", server.Addr)
		err := server.ListenAndServe()
		// http.ErrServerClosed is expected when Shutdown is called.
		if err != nil && !errors.Is(err, http.ErrServerClosed) {
			serverErrors <- err
		}
	}()

	// Block until we receive a shutdown signal or a server error.
	select {
	case err := <-serverErrors:
		log.Error().Err(err).Msg("Server error")
		return
	case <-ctx.Done():
		stop() // Reset signal behavior so a second signal exits immediately.
		log.Info().Msg("Shutdown signal received, draining requests...")
		// Give in flight requests a deadline to complete.
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		err := server.Shutdown(shutdownCtx)
		if err != nil {
			log.Error().Err(err).Msg("Graceful shutdown failed")
			return
		}
		log.Info().Msg("Server shut down gracefully")
	}

	// At this point, when function returns,
	// we make sure deferred functions run:
	// - db.Close() closes the connection pool
	// - stop() resets signal handling
}