From ee28a78edae66ccb34d1dd2bac5ce43658031688 Mon Sep 17 00:00:00 2001
From: Balakrishnan Balasubramanian <git.user@balki.me>
Date: Mon, 11 Sep 2023 20:24:10 -0400
Subject: [PATCH] Add WrapHandler to automatically call Tick on every request

---
 idle/idle.go      | 46 ++++++++++++++++++++++++++++++++++++++++++----
 idle/idle_test.go |  2 +-
 2 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/idle/idle.go b/idle/idle.go
index 105bb75..ee3d81f 100644
--- a/idle/idle.go
+++ b/idle/idle.go
@@ -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() {
 	}
 }
 
+// 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 {
-	Enter()
-	Exit()
-	Wait()
+	// 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{})
diff --git a/idle/idle_test.go b/idle/idle_test.go
index bce4c51..b172909 100644
--- a/idle/idle_test.go
+++ b/idle/idle_test.go
@@ -5,7 +5,7 @@ import (
 	"time"
 )
 
-func TestIdlerChan(t *testing.T) {
+func TestIdlerChan(_ *testing.T) {
 	i := CreateIdler(1 * time.Second)
 	<-i.Chan()
 }