From 913e33a99347c0cf2689418ea65aad1c9b372777 Mon Sep 17 00:00:00 2001
From: Muyao CHEN <muyao.chen@intersec.com>
Date: Sun, 15 Sep 2024 22:39:58 +0200
Subject: [PATCH] Context: Test context with timeout

---
 handlers.go | 58 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 main.go     | 12 +++--------
 routes.go   |  7 +++++++
 3 files changed, 68 insertions(+), 9 deletions(-)
 create mode 100644 handlers.go
 create mode 100644 routes.go

diff --git a/handlers.go b/handlers.go
new file mode 100644
index 0000000..d88bb41
--- /dev/null
+++ b/handlers.go
@@ -0,0 +1,58 @@
+package main
+
+import (
+	"context"
+	"log"
+	"net/http"
+	"time"
+
+	"git.vinchent.xyz/vinchent/go-web/framework"
+)
+
+// func FooControllerHandler(ctx *framework.Context) error {
+// 	return ctx.WriteJSON(http.StatusOK, map[string]any{
+// 		"code": 0,
+// 	})
+// }
+
+func FooControllerHandler(ctx *framework.Context) error {
+	durationCtx, cancel := context.WithTimeout(ctx.BaseContext(), time.Duration(1*time.Second))
+	defer cancel()
+
+	finish := make(chan struct{}, 1)
+	panicChan := make(chan interface{}, 1)
+
+	// some long task
+	go func() {
+		// Deal with the panic during the work
+		defer func() {
+			if p := recover(); p != nil {
+				panicChan <- p
+			}
+		}()
+		// do the business
+		time.Sleep(10 * time.Second)
+		ctx.WriteJSON(http.StatusOK, "ok")
+		finish <- struct{}{}
+	}()
+
+	select {
+	case <-panicChan:
+		// Protect ResponseWriter for concurrently writing from different
+		// goroutines if there are any
+		ctx.WriterMux().Lock()
+		defer ctx.WriterMux().Unlock()
+		log.Println("panicked")
+		ctx.WriteJSON(http.StatusInternalServerError, "panicked")
+	case <-finish:
+		log.Println("finished")
+	case <-durationCtx.Done():
+		ctx.WriterMux().Lock()
+		defer ctx.WriterMux().Unlock()
+		log.Println("Timeout")
+		ctx.WriteJSON(http.StatusInternalServerError, "time out")
+		ctx.SetHasTimeout()
+	}
+
+	return nil
+}
diff --git a/main.go b/main.go
index 99ae4c8..06f6368 100644
--- a/main.go
+++ b/main.go
@@ -1,24 +1,18 @@
 package main
 
 import (
-	"fmt"
-	"html"
 	"log"
 	"net/http"
 
 	"git.vinchent.xyz/vinchent/go-web/framework"
 )
 
-type FooHandler struct{}
-
-func (foo FooHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	fmt.Fprintf(w, "Hello, %q", html.EscapeString(r.URL.Path))
-}
-
 func main() {
+	core := framework.NewCore()
+	registerRouter(core)
 	server := &http.Server{
 		Addr:    ":8080",
-		Handler: framework.NewCore(),
+		Handler: core,
 	}
 
 	if err := server.ListenAndServe(); err != nil {
diff --git a/routes.go b/routes.go
new file mode 100644
index 0000000..659d8b1
--- /dev/null
+++ b/routes.go
@@ -0,0 +1,7 @@
+package main
+
+import "git.vinchent.xyz/vinchent/go-web/framework"
+
+func registerRouter(core *framework.Core) {
+	core.Get("foo", FooControllerHandler)
+}