From eecd5ba059219a23673f4ed83984b096501a9fee Mon Sep 17 00:00:00 2001
From: Neur0toxine <pashok9825@gmail.com>
Date: Thu, 19 Sep 2024 18:05:16 +0300
Subject: [PATCH] stream id middleware

---
 core/logger/buffer_logger_test.go |  2 +-
 core/logger/gin.go                | 11 +++++++-
 core/logger/gin_test.go           | 23 +++++++++-------
 core/logger/stream_id.go          | 44 +++++++++++++++++++++++++++++++
 core/logger/stream_id_test.go     | 19 +++++++++++++
 5 files changed, 88 insertions(+), 11 deletions(-)
 create mode 100644 core/logger/stream_id.go
 create mode 100644 core/logger/stream_id_test.go

diff --git a/core/logger/buffer_logger_test.go b/core/logger/buffer_logger_test.go
index 0c6115e..d1a36ed 100644
--- a/core/logger/buffer_logger_test.go
+++ b/core/logger/buffer_logger_test.go
@@ -17,7 +17,7 @@ import (
 type logRecord struct {
 	LevelName  string                 `json:"level_name"`
 	DateTime   null.Time              `json:"datetime"`
-	Caller     string                 `json:"caller"`
+	StreamID   string                 `json:"streamId"`
 	Message    string                 `json:"message"`
 	Handler    string                 `json:"handler,omitempty"`
 	Connection string                 `json:"connection,omitempty"`
diff --git a/core/logger/gin.go b/core/logger/gin.go
index 5ecf3b2..9c6ded0 100644
--- a/core/logger/gin.go
+++ b/core/logger/gin.go
@@ -7,7 +7,7 @@ import (
 	"go.uber.org/zap"
 )
 
-// GinMiddleware will construct Gin middleware which will log requests.
+// GinMiddleware will construct Gin middleware which will log requests and provide logger with unique request ID.
 func GinMiddleware(log Logger) gin.HandlerFunc {
 	return func(c *gin.Context) {
 		// Start timer
@@ -15,6 +15,11 @@ func GinMiddleware(log Logger) gin.HandlerFunc {
 		path := c.Request.URL.Path
 		raw := c.Request.URL.RawQuery
 
+		streamID := generateStreamID()
+		log := log.With(StreamID(streamID))
+		c.Set(StreamIDAttr, streamID)
+		c.Set("logger", log)
+
 		// Process request
 		c.Next()
 
@@ -35,3 +40,7 @@ func GinMiddleware(log Logger) gin.HandlerFunc {
 		)
 	}
 }
+
+func MustGet(c *gin.Context) Logger {
+	return c.MustGet("logger").(Logger)
+}
diff --git a/core/logger/gin_test.go b/core/logger/gin_test.go
index 3cf6cf7..6558eda 100644
--- a/core/logger/gin_test.go
+++ b/core/logger/gin_test.go
@@ -16,6 +16,8 @@ func TestGinMiddleware(t *testing.T) {
 	r := gin.New()
 	r.Use(GinMiddleware(log))
 	r.GET("/mine", func(c *gin.Context) {
+		log := MustGet(c)
+		log.Info("some very important message")
 		c.JSON(http.StatusOK, gin.H{})
 	})
 	r.ServeHTTP(rr, httptest.NewRequest(http.MethodGet, "/mine", nil))
@@ -23,16 +25,19 @@ func TestGinMiddleware(t *testing.T) {
 	require.Equal(t, http.StatusOK, rr.Code)
 	items, err := newJSONBufferedLogger(log).ScanAll()
 	require.NoError(t, err)
-	require.Len(t, items, 1)
-	require.NotEmpty(t, items[0].Context)
-	assert.NotEmpty(t, items[0].Context["startTime"])
-	assert.NotEmpty(t, items[0].Context["endTime"])
+	require.Len(t, items, 2)
+	require.NotEmpty(t, items[0].StreamID)
+	require.NotEmpty(t, items[1].StreamID)
+	assert.Equal(t, "some very important message", items[0].Message)
+	require.NotEmpty(t, items[1].Context)
+	assert.NotEmpty(t, items[1].Context["startTime"])
+	assert.NotEmpty(t, items[1].Context["endTime"])
 	assert.True(t, func() bool {
-		_, ok := items[0].Context["latency"]
+		_, ok := items[1].Context["latency"]
 		return ok
 	}())
-	assert.NotEmpty(t, items[0].Context["remoteAddress"])
-	assert.NotEmpty(t, items[0].Context[HTTPMethodAttr])
-	assert.NotEmpty(t, items[0].Context["path"])
-	assert.NotEmpty(t, items[0].Context["bodySize"])
+	assert.NotEmpty(t, items[1].Context["remoteAddress"])
+	assert.NotEmpty(t, items[1].Context[HTTPMethodAttr])
+	assert.NotEmpty(t, items[1].Context["path"])
+	assert.NotEmpty(t, items[1].Context["bodySize"])
 }
diff --git a/core/logger/stream_id.go b/core/logger/stream_id.go
new file mode 100644
index 0000000..2ba8e45
--- /dev/null
+++ b/core/logger/stream_id.go
@@ -0,0 +1,44 @@
+package logger
+
+import (
+	"encoding/binary"
+	"sync/atomic"
+	"time"
+	"unsafe"
+)
+
+var streamIDCounter = uint64(time.Now().Unix())
+
+func hexEncode(dst *[32]byte, src *[8]byte, cnt uint64) int {
+	const hexDigits = "0123456789abcdef"
+	for i, b := range src[:] {
+		dst[i*2] = hexDigits[b>>4]
+		dst[i*2+1] = hexDigits[b&0x0F]
+	}
+	idx := 16
+	for cnt > 0 {
+		dst[idx] = hexDigits[cnt&0xF]
+		cnt >>= 4
+		idx++
+	}
+	for i, j := 16, idx-1; i < j; i, j = i+1, j-1 {
+		dst[i], dst[j] = dst[j], dst[i]
+	}
+	return idx
+}
+
+func bytesToString(b *[32]byte, length int) string {
+	return unsafe.String(&b[0], length)
+}
+
+func generateStreamID() string {
+	var id [8]byte
+	var hexID [32]byte
+
+	cnt := atomic.AddUint64(&streamIDCounter, 1)
+	timestamp := time.Now().UnixNano()
+	binary.BigEndian.PutUint64(id[:], uint64(timestamp))
+	length := hexEncode(&hexID, &id, cnt)
+
+	return bytesToString(&hexID, length)
+}
diff --git a/core/logger/stream_id_test.go b/core/logger/stream_id_test.go
new file mode 100644
index 0000000..19f0f8e
--- /dev/null
+++ b/core/logger/stream_id_test.go
@@ -0,0 +1,19 @@
+package logger
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestGenerateStreamID(t *testing.T) {
+	id1 := generateStreamID()
+	id2 := generateStreamID()
+
+	assert.NotEqual(t, id1, id2)
+}
+
+func BenchmarkGenerateStreamID(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		_ = generateStreamID()
+	}
+}