main.go

Documentation: github.com/go-chi/chi/_examples/graceful

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