main.go

package main

import (
	"context"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/signal"
	"sync/atomic"
	"syscall"
	"time"
)

var startup = time.Now()
var delay = 30 * time.Second
var drain int32

func main() {
	if del := os.Getenv("STARTUP_DELAY"); del != "" {
		var err error
		delay, err = time.ParseDuration(del)
		if err != nil {
			log.Fatal(err)
		}
	}

	srv := &http.Server{Addr: ":8080"}
	go func() {
		ch := make(chan os.Signal)
		signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
		sig := <-ch
		log.Printf("Got signal %v, draining requests", sig)
		atomic.StoreInt32(&drain, 1)
		ctx, _ := context.WithTimeout(context.Background(), 10*time.Second)
		if err := srv.Shutdown(ctx); err != nil {
			log.Fatalf("Shutdown: %v", err)
		}
		log.Println("Shutdown successful")
	}()

	http.HandleFunc("/healthz", health)
	http.HandleFunc("/", serve)
	log.Fatal(srv.ListenAndServe())
}

func serve(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "since startup: %s\n", time.Now().Sub(startup))

	for _, e := range os.Environ() {
		fmt.Fprintln(w, e)
	}
}

func health(w http.ResponseWriter, r *http.Request) {
	if atomic.LoadInt32(&drain) == 1 {
		http.Error(w, "draining", http.StatusServiceUnavailable)
		return
	}
	if time.Now().After(startup.Add(delay)) {
		io.WriteString(w, "ok")
		return
	}
	http.Error(w, "starting up", http.StatusServiceUnavailable)
}