From 8c3cbffd6cd34f975621987fb4dc431e1fe33406 Mon Sep 17 00:00:00 2001
From: vearne <asdwoshiaotian@gmail.com>
Date: Tue, 28 May 2024 15:09:31 +0800
Subject: [PATCH] use atomic.Bool

---
 example/file_server/file_server.go |  4 ++--
 timeout.go                         |  4 ++--
 writer.go                          | 11 ++++++-----
 3 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/example/file_server/file_server.go b/example/file_server/file_server.go
index 445d355..264ffbb 100644
--- 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
index 973c254..6f675a9 100644
--- 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
index 18b08e7..d048da2 100644
--- 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
 }