X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/4ed4b6554535849341673efb7f80392dd5fba946..c93eb8717c79922f3e68b13405b5250bb1345d72:/sdk/go/httpserver/id_generator.go

diff --git a/sdk/go/httpserver/id_generator.go b/sdk/go/httpserver/id_generator.go
index c2830f7865..d2c3a41f21 100644
--- a/sdk/go/httpserver/id_generator.go
+++ b/sdk/go/httpserver/id_generator.go
@@ -1,6 +1,12 @@
+// Copyright (C) The Arvados Authors. All rights reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+
 package httpserver
 
 import (
+	"math/rand"
+	"net/http"
 	"strconv"
 	"sync"
 	"time"
@@ -13,19 +19,34 @@ 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("X-Request-Id") == "" {
+			req.Header.Set("X-Request-Id", gen.Next())
+		}
+		h.ServeHTTP(w, req)
+	})
 }