X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/0561bd0c3c07257fd58ded6c7cfa5feeae97af57..8d6c35e4dd68baec1b58d566d32d89feb651deac:/sdk/go/httpserver/id_generator.go

diff --git a/sdk/go/httpserver/id_generator.go b/sdk/go/httpserver/id_generator.go
index 18fd91f560..5b661ae9f7 100644
--- a/sdk/go/httpserver/id_generator.go
+++ b/sdk/go/httpserver/id_generator.go
@@ -5,11 +5,17 @@
 package httpserver
 
 import (
+	"math/rand"
+	"net/http"
 	"strconv"
 	"sync"
 	"time"
 )
 
+const (
+	HeaderRequestID = "X-Request-Id"
+)
+
 // IDGenerator generates alphanumeric strings suitable for use as
 // unique IDs (a given IDGenerator will never return the same ID
 // twice).
@@ -17,19 +23,38 @@ type IDGenerator struct {
 	// Prefix is prepended to each returned ID.
 	Prefix string
 
-	lastID int64
-	mtx    sync.Mutex
+	mtx sync.Mutex
+	src rand.Source
 }
 
 // Next returns a new ID string. It is safe to call Next from multiple
 // goroutines.
 func (g *IDGenerator) Next() string {
-	id := time.Now().UnixNano()
 	g.mtx.Lock()
-	if id <= g.lastID {
-		id = g.lastID + 1
+	defer g.mtx.Unlock()
+	if g.src == nil {
+		g.src = rand.NewSource(time.Now().UnixNano())
+	}
+	a, b := g.src.Int63(), g.src.Int63()
+	id := strconv.FormatInt(a, 36) + strconv.FormatInt(b, 36)
+	for len(id) > 20 {
+		id = id[:20]
 	}
-	g.lastID = id
-	g.mtx.Unlock()
-	return g.Prefix + strconv.FormatInt(id, 36)
+	return g.Prefix + id
+}
+
+// AddRequestIDs wraps an http.Handler, adding an X-Request-Id header
+// to each request that doesn't already have one.
+func AddRequestIDs(h http.Handler) http.Handler {
+	gen := &IDGenerator{Prefix: "req-"}
+	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+		if req.Header.Get(HeaderRequestID) == "" {
+			if req.Header == nil {
+				req.Header = http.Header{}
+			}
+			req.Header.Set(HeaderRequestID, gen.Next())
+		}
+		w.Header().Set("X-Request-Id", req.Header.Get("X-Request-Id"))
+		h.ServeHTTP(w, req)
+	})
 }