Merge pull request #38 from vearne/fix/bool

use atomic.Bool
vearne · May 28, 2024 · a82a1f8 · a82a1f8
2 parents 74f243a + 8c3cbff
commit a82a1f8
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 9 deletions.
diff --git a/example/file_server/file_server.go b/example/file_server/file_server.go
@@ -18,8 +18,8 @@ func main() {
 	log.Fatal(router.Run(":8080"))
 }
 
-// mkdir -p /tmp/foo
-// echo "a" >> /tmp/foo/a
+// mkdir -p /tmp/static
+// echo "a" >> /tmp/static/a
 
 // test case1:
 // curl -I http://localhost:8080/static/a

diff --git a/timeout.go b/timeout.go
@@ -80,7 +80,7 @@ func Timeout(opts ...Option) gin.HandlerFunc {
 			tw.mu.Lock()
 			defer tw.mu.Unlock()
 
-			tw.timedOut = true
+			tw.timedOut.Store(true)
 			tw.ResponseWriter.WriteHeader(tw.ErrorHttpCode)
 
 			n, err = tw.ResponseWriter.Write(encodeBytes(tw.DefaultMsg))
@@ -107,7 +107,7 @@ func Timeout(opts ...Option) gin.HandlerFunc {
 				dst[k] = vv
 			}
 
-			if !tw.wroteHeader {
+			if !tw.wroteHeader.Load() {
 				tw.code = c.Writer.Status()
 			}
 

diff --git a/writer.go b/writer.go
@@ -5,6 +5,7 @@ import (
 	"github.com/gin-gonic/gin"
 	"net/http"
 	"sync"
+	"sync/atomic"
 )
 
 type TimeoutWriter struct {
@@ -17,15 +18,15 @@ type TimeoutWriter struct {
 
 	code        int
 	mu          sync.Mutex
-	timedOut    bool
-	wroteHeader bool
+	timedOut    atomic.Bool
+	wroteHeader atomic.Bool
 	size        int
 }
 
 func (tw *TimeoutWriter) Write(b []byte) (int, error) {
 	tw.mu.Lock()
 	defer tw.mu.Unlock()
-	if tw.timedOut {
+	if tw.timedOut.Load() {
 		return 0, nil
 	}
 	tw.size += len(b)
@@ -35,14 +36,14 @@ func (tw *TimeoutWriter) Write(b []byte) (int, error) {
 func (tw *TimeoutWriter) WriteHeader(code int) {
 	tw.mu.Lock()
 	defer tw.mu.Unlock()
-	if tw.timedOut {
+	if tw.timedOut.Load() {
 		return
 	}
 	tw.writeHeader(code)
 }
 
 func (tw *TimeoutWriter) writeHeader(code int) {
-	tw.wroteHeader = true
+	tw.wroteHeader.Store(true)
 	tw.code = code
 }