Add WrapHandler to automatically call Tick on every request

This commit is contained in:
Balakrishnan Balasubramanian 2023-09-11 20:24:10 -04:00
parent 8ca82581fe
commit ee28a78eda
2 changed files with 43 additions and 5 deletions

View File

@ -1,16 +1,20 @@
// Package idle helps to gracefully shutdown idle servers // Package idle helps to gracefully shutdown idle (typically http) servers
package idle package idle
import ( import (
"fmt" "fmt"
"net/http"
"sync/atomic" "sync/atomic"
"time" "time"
) )
var ( var (
// For simple servers without backgroud jobs, global singleton for simpler API
// Enter/Exit worn't work for global idler as Enter may be called before Wait, use CreateIdler in those cases
gIdler atomic.Pointer[idler] gIdler atomic.Pointer[idler]
) )
// Wait waits till the server is idle and returns. i.e. no Ticks in last <timeout> duration
func Wait(timeout time.Duration) error { func Wait(timeout time.Duration) error {
i := CreateIdler(timeout).(*idler) i := CreateIdler(timeout).(*idler)
ok := gIdler.CompareAndSwap(nil, i) ok := gIdler.CompareAndSwap(nil, i)
@ -21,6 +25,7 @@ func Wait(timeout time.Duration) error {
return nil return nil
} }
// Tick records the current time. This will make the server not idle until next Tick or timeout
func Tick() { func Tick() {
i := gIdler.Load() i := gIdler.Load()
if i != nil { if i != nil {
@ -28,11 +33,43 @@ func Tick() {
} }
} }
type Idler interface { // WrapHandler calls Tick() before processing passing request to http.Handler
Enter() func WrapHandler(h http.Handler) http.Handler {
Exit() if h == nil {
Wait() h = http.DefaultServeMux
}
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
Tick() Tick()
h.ServeHTTP(w, r)
})
}
// WrapIdlerHandler calls idler.Tick() before processing passing request to http.Handler
func WrapIdlerHandler(i Idler, h http.Handler) http.Handler {
if h == nil {
h = http.DefaultServeMux
}
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
i.Tick()
h.ServeHTTP(w, r)
})
}
// Idler helps manage idle servers
type Idler interface {
// Tick records the current time. This will make the server not idle until next Tick or timeout
Tick()
// Wait waits till the server is idle and returns. i.e. no Ticks in last <timeout> duration
Wait()
// For long running background jobs, use Enter to record start time. Wait will not return while there are active jobs running
Enter()
// Exit records end of a background job
Exit()
// Get the channel to wait yourself
Chan() <-chan struct{} Chan() <-chan struct{}
} }
@ -51,6 +88,7 @@ func (i *idler) Exit() {
i.active.Add(-1) i.active.Add(-1)
} }
// CreateIdler creates an Idler with given timeout
func CreateIdler(timeout time.Duration) Idler { func CreateIdler(timeout time.Duration) Idler {
i := &idler{} i := &idler{}
i.c = make(chan struct{}) i.c = make(chan struct{})

View File

@ -5,7 +5,7 @@ import (
"time" "time"
) )
func TestIdlerChan(t *testing.T) { func TestIdlerChan(_ *testing.T) {
i := CreateIdler(1 * time.Second) i := CreateIdler(1 * time.Second)
<-i.Chan() <-i.Chan()
} }