Merge branch 'thehyve/fix-crunch-documentation' Fix a typo in Crunch Dispatch install...

[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 c2830f7865da86541294ca806090ea73a6be4955..6452136d85eede6896f1dca1648e00b4ba6ae8e7 100644 (file)
--- 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,37 @@ 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") == "" {
+                       if req.Header == nil {
+                               req.Header = http.Header{}
+                       }
+                       req.Header.Set("X-Request-Id", gen.Next())
+               }
+               h.ServeHTTP(w, req)
+       })
 }