Skip to content

Commit

Permalink
stream id middleware
Browse files Browse the repository at this point in the history
  • Loading branch information
Neur0toxine committed Sep 19, 2024
1 parent 4dcb31f commit eecd5ba
Show file tree
Hide file tree
Showing 5 changed files with 88 additions and 11 deletions.
2 changes: 1 addition & 1 deletion core/logger/buffer_logger_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"`
Expand Down
11 changes: 10 additions & 1 deletion core/logger/gin.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,14 +7,19 @@ 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
start := time.Now()
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()

Expand All @@ -35,3 +40,7 @@ func GinMiddleware(log Logger) gin.HandlerFunc {
)
}
}

func MustGet(c *gin.Context) Logger {
return c.MustGet("logger").(Logger)
}
23 changes: 14 additions & 9 deletions core/logger/gin_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,23 +16,28 @@ 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))

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"])
}
44 changes: 44 additions & 0 deletions core/logger/stream_id.go
Original file line number Diff line number Diff line change
@@ -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)
}
19 changes: 19 additions & 0 deletions core/logger/stream_id_test.go
Original file line number Diff line number Diff line change
@@ -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()
}
}

0 comments on commit eecd5ba

Please sign in to comment.