diff --git a/server.go b/server.go index 0673ca3..b17cae4 100644 --- a/server.go +++ b/server.go @@ -359,7 +359,9 @@ func (ctx *RequestCtx) Hijack(handler HijackHandler) { // This functionality may be useful for passing arbitrary values between // functions involved in request processing. // -// All the values stored in ctx are deleted after returning from RequestHandler. +// All the values are removed from ctx after returning from the top +// RequestHandler. Additionally, Close method is called on each value +// implementing io.Closer before removing the value from ctx. func (ctx *RequestCtx) SetUserValue(key string, value interface{}) { ctx.userValues.Set(key, value) } diff --git a/userdata.go b/userdata.go index ecc343e..a99ec5d 100644 --- a/userdata.go +++ b/userdata.go @@ -1,5 +1,9 @@ package fasthttp +import ( + "io" +) + type userDataKV struct { key []byte value interface{} @@ -55,5 +59,13 @@ func (d *userData) GetBytes(key []byte) interface{} { } func (d *userData) Reset() { + args := *d + n := len(args) + for i := 0; i < n; i++ { + v := args[i].value + if vc, ok := v.(io.Closer); ok { + vc.Close() + } + } *d = (*d)[:0] } diff --git a/userdata_test.go b/userdata_test.go index 995bf81..ce14204 100644 --- a/userdata_test.go +++ b/userdata_test.go @@ -39,3 +39,36 @@ func testUserDataGet(t *testing.T, u *userData, key []byte, value interface{}) { t.Fatalf("unexpected value for key=%q: %d. Expecting %d", key, v, value) } } + +func TestUserDataValueClose(t *testing.T) { + var u userData + + closeCalls := 0 + + // store values implementing io.Closer + for i := 0; i < 5; i++ { + key := fmt.Sprintf("key_%d", i) + u.Set(key, &closerValue{&closeCalls}) + } + + // store values without io.Closer + for i := 0; i < 10; i++ { + key := fmt.Sprintf("key_noclose_%d", i) + u.Set(key, i) + } + + u.Reset() + + if closeCalls != 5 { + t.Fatalf("unexpected number of Close calls: %d. Expecting 10", closeCalls) + } +} + +type closerValue struct { + closeCalls *int +} + +func (cv *closerValue) Close() error { + (*cv.closeCalls)++ + return nil +}