sdk/go/httpserver/id_generator.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: Apache-2.0
   4
   5 package httpserver
   6
   7 import (
   8         "math/rand"
   9         "net/http"
  10         "strconv"
  11         "sync"
  12         "time"
  13 )
  14
  15 const (
  16         HeaderRequestID = "X-Request-Id"
  17 )
  18
  19 // IDGenerator generates alphanumeric strings suitable for use as
  20 // unique IDs (a given IDGenerator will never return the same ID
  21 // twice).
  22 type IDGenerator struct {
  23         // Prefix is prepended to each returned ID.
  24         Prefix string
  25
  26         mtx sync.Mutex
  27         src rand.Source
  28 }
  29
  30 // Next returns a new ID string. It is safe to call Next from multiple
  31 // goroutines.
  32 func (g *IDGenerator) Next() string {
  33         g.mtx.Lock()
  34         defer g.mtx.Unlock()
  35         if g.src == nil {
  36                 g.src = rand.NewSource(time.Now().UnixNano())
  37         }
  38         a, b := g.src.Int63(), g.src.Int63()
  39         id := strconv.FormatInt(a, 36) + strconv.FormatInt(b, 36)
  40         for len(id) > 20 {
  41                 id = id[:20]
  42         }
  43         return g.Prefix + id
  44 }
  45
  46 // AddRequestIDs wraps an http.Handler, adding an X-Request-Id header
  47 // to each request that doesn't already have one.
  48 func AddRequestIDs(h http.Handler) http.Handler {
  49         gen := &IDGenerator{Prefix: "req-"}
  50         return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
  51                 if req.Header.Get(HeaderRequestID) == "" {
  52                         if req.Header == nil {
  53                                 req.Header = http.Header{}
  54                         }
  55                         req.Header.Set(HeaderRequestID, gen.Next())
  56                 }
  57                 w.Header().Set("X-Request-Id", req.Header.Get("X-Request-Id"))
  58                 h.ServeHTTP(w, req)
  59         })
  60 }