From 07cd03a67423ed3bfd7c0befd662c23f4b2e8807 Mon Sep 17 00:00:00 2001
From: "Gerasimos (Makis) Maropoulos" <kataras2006@hotmail.com>
Date: Mon, 18 May 2020 00:10:09 +0300
Subject: [PATCH] make context.OnClose's callback wrapped by sync.Once to make
sure that the callback is only called once
AQQkADAwATZiZmYAZC05YzI0LTdmOTAtMDACLTAwCgAQAHsNZaaCfV1BmxBvtU
Former-commit-id: 9449270ccf276ea1bdf51fbdde03b81223290e2a
---
context/context.go | 22 +++++++++++++++++-----
1 file changed, 17 insertions(+), 5 deletions(-)
diff --git a/context/context.go b/context/context.go
index ade8532b..70d708b8 100644
--- a/context/context.go
+++ b/context/context.go
@@ -21,6 +21,7 @@ import (
"regexp"
"strconv"
"strings"
+ "sync"
"time"
"unsafe"
@@ -307,6 +308,7 @@ type Context interface {
// OnClose registers the callback function "cb" to the underline connection closing event using the `Context#OnConnectionClose`
// and also in the end of the request handler using the `ResponseWriter#SetBeforeFlush`.
// Note that you can register only one callback for the entire request handler chain/per route.
+ // Note that the "cb" will only be called once.
//
// Look the `Context#OnConnectionClose` and `ResponseWriter#SetBeforeFlush` for more.
OnClose(cb func())
@@ -1633,6 +1635,10 @@ func (ctx *context) StopWithProblem(statusCode int, problem Problem) {
//
// Look the `ResponseWriter#CloseNotifier` for more.
func (ctx *context) OnConnectionClose(cb func()) bool {
+ if cb == nil {
+ return false
+ }
+
// Note that `ctx.ResponseWriter().CloseNotify()` can already do the same
// but it returns a channel which will never fire if it the protocol version is not compatible,
// here we don't want to allocate an empty channel, just skip it.
@@ -1644,9 +1650,7 @@ func (ctx *context) OnConnectionClose(cb func()) bool {
notify := notifier.CloseNotify()
go func() {
<-notify
- if cb != nil {
- cb()
- }
+ cb()
}()
return true
@@ -1656,14 +1660,22 @@ func (ctx *context) OnConnectionClose(cb func()) bool {
// and also in the end of the request handler using the `ResponseWriter#SetBeforeFlush`.
// Note that you can register only one callback for the entire request handler chain/per route.
//
+// Note that the "cb" will only be called once.
+//
// Look the `Context#OnConnectionClose` and `ResponseWriter#SetBeforeFlush` for more.
func (ctx *context) OnClose(cb func()) {
if cb == nil {
return
}
+ once := new(sync.Once)
+
+ callOnce := func() {
+ once.Do(cb)
+ }
+
// Register the on underline connection close handler first.
- ctx.OnConnectionClose(cb)
+ ctx.OnConnectionClose(callOnce)
// Author's notes:
// This is fired on `ctx.ResponseWriter().FlushResponse()` which is fired by the framework automatically, internally, on the end of request handler(s),
@@ -1680,7 +1692,7 @@ func (ctx *context) OnClose(cb func()) {
// return
// }
- ctx.writer.SetBeforeFlush(cb)
+ ctx.writer.SetBeforeFlush(callOnce)
}
// +------------------------------------------------------------+