From ae61c328603f5b6f8b9f933ccbb9e630ece13dc0 Mon Sep 17 00:00:00 2001
From: Muhammed Efe Cetin <efectn@protonmail.com>
Date: Sat, 29 Jun 2024 01:15:42 +0300
Subject: [PATCH] WIP: export buildTree method

---
 app.go         |  2 +-
 router.go      | 22 +++++++++++++++++++--
 router_test.go | 53 ++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 74 insertions(+), 3 deletions(-)

diff --git a/app.go b/app.go
index 01941edd..d0ab6981 100644
--- a/app.go
+++ b/app.go
@@ -86,7 +86,7 @@ type Error struct {
 
 // App denotes the Fiber application.
 type App struct {
-	mutex sync.Mutex
+	mutex sync.RWMutex
 	// Route stack divided by HTTP methods
 	stack [][]*Route
 	// Route stack divided by HTTP methods and route prefixes
diff --git a/router.go b/router.go
index eae9adef..cf36dede 100644
--- a/router.go
+++ b/router.go
@@ -146,12 +146,15 @@ func (app *App) nextCustom(c CustomCtx) (bool, error) { //nolint: unparam // boo
 }
 
 func (app *App) next(c *DefaultCtx) (bool, error) {
+	app.mutex.Lock()
+
 	// Get stack length
 	tree, ok := app.treeStack[c.methodINT][c.treePath]
 	if !ok {
 		tree = app.treeStack[c.methodINT][""]
 	}
 	lenTree := len(tree) - 1
+	app.mutex.Unlock()
 
 	// Loop over the route stack starting from previous index
 	for c.indexRoute < lenTree {
@@ -375,6 +378,11 @@ func (app *App) register(methods []string, pathRaw string, group *Group, handler
 }
 
 func (app *App) addRoute(method string, route *Route, isMounted ...bool) {
+	app.mutex.Lock()
+	defer app.mutex.Unlock()
+
+	fmt.Printf("addRoute: method: %s, route: %v, isMounted: %v\n", method, route, isMounted)
+
 	// Check mounted routes
 	var mounted bool
 	if len(isMounted) > 0 {
@@ -400,12 +408,12 @@ func (app *App) addRoute(method string, route *Route, isMounted ...bool) {
 
 	// Execute onRoute hooks & change latestRoute if not adding mounted route
 	if !mounted {
-		app.mutex.Lock()
+		//app.mutex.Lock()
 		app.latestRoute = route
 		if err := app.hooks.executeOnRouteHooks(*route); err != nil {
 			panic(err)
 		}
-		app.mutex.Unlock()
+		//app.mutex.Unlock()
 	}
 }
 
@@ -446,3 +454,13 @@ func (app *App) buildTree() *App {
 
 	return app
 }
+
+// BuildTree rebuilds the prefix tree from the previously registered routes.
+// This method is useful when you want to register routes dynamically after the app has started.
+// It is not recommended to use this method on produtcion environments because rebuilding the tree is performance-intensive.
+func (app *App) BuildTree() *App {
+	app.mutex.Lock()
+	defer app.mutex.Unlock()
+
+	return app.buildTree()
+}
diff --git a/router_test.go b/router_test.go
index 57ce9209..300614b2 100644
--- a/router_test.go
+++ b/router_test.go
@@ -11,6 +11,9 @@ import (
 	"io"
 	"net/http/httptest"
 	"os"
+	"strconv"
+	"sync"
+	"sync/atomic"
 	"testing"
 
 	"github.com/gofiber/utils/v2"
@@ -368,6 +371,56 @@ func Test_Router_NotFound_HTML_Inject(t *testing.T) {
 	require.Equal(t, "Cannot DELETE /does/not/exist&lt;script&gt;alert(&#39;foo&#39;);&lt;/script&gt;", string(c.Response.Body()))
 }
 
+func Test_App_BuildTree_Concurrent(t *testing.T) {
+	t.Parallel()
+	app := New()
+
+	var counter atomic.Uint64
+	var routerMu sync.Mutex
+	app.Get("/test", func(c Ctx) error {
+		routerMu.Lock()
+
+		cStr := strconv.FormatUint(counter.Load(), 10)
+		counter.Add(1)
+
+		app.Get("/test"+cStr, func(c Ctx) error {
+			return c.SendString("test" + cStr)
+		})
+
+		routerMu.Unlock()
+
+		app.BuildTree()
+
+		return c.SendStatus(StatusOK)
+	})
+
+	var wg sync.WaitGroup
+	wg.Add(4)
+	for i := 0; i < 4; i++ {
+		go func() {
+			defer wg.Done()
+			req := httptest.NewRequest(MethodGet, "/test", nil)
+			resp, err := app.Test(req)
+
+			require.NoError(t, err)
+			require.Equal(t, StatusOK, resp.StatusCode)
+		}()
+	}
+	wg.Wait()
+
+	for i := 0; i < 4; i++ {
+		req := httptest.NewRequest(MethodGet, "/test"+strconv.FormatInt(int64(i), 10), nil)
+		resp, err := app.Test(req)
+
+		require.NoError(t, err)
+		require.Equal(t, StatusOK, resp.StatusCode)
+
+		body, err := io.ReadAll(resp.Body)
+		require.NoError(t, err)
+		require.Equal(t, "test"+strconv.Itoa(i), app.getString(body))
+	}
+}
+
 //////////////////////////////////////////////
 ///////////////// BENCHMARKS /////////////////
 //////////////////////////////////////////////