1package main
2
3import (
4	"context"
5	"fmt"
6	"net"
7	"net/http"
8	"os"
9	"os/signal"
10	"time"
11
12	"github.com/go-chi/chi"
13	"github.com/go-chi/chi/middleware"
14	"github.com/go-chi/valve"
15)
16
17func main() {
18
19	// Our graceful valve shut-off package to manage code preemption and
20	// shutdown signaling.
21	valv := valve.New()
22	baseCtx := valv.Context()
23
24	// Example of a long running background worker thing..
25	go func(ctx context.Context) {
26		for {
27			<-time.After(1 * time.Second)
28
29			func() {
30				valve.Lever(ctx).Open()
31				defer valve.Lever(ctx).Close()
32
33				// actual code doing stuff..
34				fmt.Println("tick..")
35				time.Sleep(2 * time.Second)
36				// end-logic
37
38				// signal control..
39				select {
40				case <-valve.Lever(ctx).Stop():
41					fmt.Println("valve is closed")
42					return
43
44				case <-ctx.Done():
45					fmt.Println("context is cancelled, go home.")
46					return
47				default:
48				}
49			}()
50
51		}
52	}(baseCtx)
53
54	// HTTP service running in this program as well. The valve context is set
55	// as a base context on the server listener at the point where we instantiate
56	// the server - look lower.
57	r := chi.NewRouter()
58	r.Use(middleware.RequestID)
59	r.Use(middleware.Logger)
60
61	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
62		w.Write([]byte("sup"))
63	})
64
65	r.Get("/slow", func(w http.ResponseWriter, r *http.Request) {
66
67		valve.Lever(r.Context()).Open()
68		defer valve.Lever(r.Context()).Close()
69
70		select {
71		case <-valve.Lever(r.Context()).Stop():
72			fmt.Println("valve is closed. finish up..")
73
74		case <-time.After(5 * time.Second):
75			// The above channel simulates some hard work.
76			// We want this handler to complete successfully during a shutdown signal,
77			// so consider the work here as some background routine to fetch a long running
78			// search query to find as many results as possible, but, instead we cut it short
79			// and respond with what we have so far. How a shutdown is handled is entirely
80			// up to the developer, as some code blocks are preemptable, and others are not.
81			time.Sleep(5 * time.Second)
82		}
83
84		w.Write([]byte(fmt.Sprintf("all done.\n")))
85	})
86
87	srv := http.Server{Addr: ":3333", Handler: r}
88	srv.BaseContext = func(_ net.Listener) context.Context {
89		return baseCtx
90	}
91
92	c := make(chan os.Signal, 1)
93	signal.Notify(c, os.Interrupt)
94	go func() {
95		for range c {
96			// sig is a ^C, handle it
97			fmt.Println("shutting down..")
98
99			// first valv
100			valv.Shutdown(20 * time.Second)
101
102			// create context with timeout
103			ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
104			defer cancel()
105
106			// start http shutdown
107			srv.Shutdown(ctx)
108
109			// verify, in worst case call cancel via defer
110			select {
111			case <-time.After(21 * time.Second):
112				fmt.Println("not all connections done")
113			case <-ctx.Done():
114
115			}
116		}
117	}()
118	srv.ListenAndServe()
119}
120