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
import (
"fmt"
"net/http"
"sync/atomic"
"time"
)
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]
)
// Wait waits till the server is idle and returns. i.e. no Ticks in last <timeout> duration
func Wait(timeout time.Duration) error {
i := CreateIdler(timeout).(*idler)
ok := gIdler.CompareAndSwap(nil, i)
@ -21,6 +25,7 @@ func Wait(timeout time.Duration) error {
return nil
}
// Tick records the current time. This will make the server not idle until next Tick or timeout
func Tick() {
i := gIdler.Load()
if i != nil {
@ -28,11 +33,43 @@ func Tick() {
}
}
type Idler interface {
Enter()
Exit()
Wait()
// WrapHandler calls Tick() before processing passing request to http.Handler
func WrapHandler(h http.Handler) http.Handler {
if h == nil {
h = http.DefaultServeMux
}
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
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{}
}
@ -51,6 +88,7 @@ func (i *idler) Exit() {
i.active.Add(-1)
}
// CreateIdler creates an Idler with given timeout
func CreateIdler(timeout time.Duration) Idler {
i := &idler{}
i.c = make(chan struct{})

View File

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