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