19205: Fix flaky test.

[arvados.git] / sdk / go / httpserver / id_generator.go

diff --git a/sdk/go/httpserver/id_generator.go b/sdk/go/httpserver/id_generator.go
index 18fd91f560227ee1c40fec2bce64126bc0c3b827..5b661ae9f71ac4984c3a22a6dfe8a58bbffbdb45 100644 (file)
--- 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)
+       })
 }