diff --git a/contrib/middleware/timeout.go b/contrib/middleware/timeout.go new file mode 100644 index 0000000000000000000000000000000000000000..99d7f69122aad9ae73efbecaf2c3851e7823eb49 --- /dev/null +++ b/contrib/middleware/timeout.go @@ -0,0 +1,26 @@ +package middleware + +import ( + "context" + "time" + + "gfx.cafe/open/jrpc/pkg/jsonrpc" +) + +func Timeout(dur time.Duration) func(h jsonrpc.Handler) jsonrpc.Handler { + return func(next jsonrpc.Handler) jsonrpc.Handler { + return jsonrpc.HandlerFunc(func(w jsonrpc.ResponseWriter, r *jsonrpc.Request) { + done := make(chan struct{}) + ctx, cn := context.WithTimeout(r.Context(), dur) + r = r.WithContext(ctx) + go func() { + next.ServeRPC(w, r) + close(done) + cn() + }() + <-ctx.Done() + cn() + <-done + }) + } +}