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 // IDGenerator generates alphanumeric strings suitable for use as
  16 // unique IDs (a given IDGenerator will never return the same ID
  17 // twice).
  18 type IDGenerator struct {
  19         // Prefix is prepended to each returned ID.
  20         Prefix string
  21
  22         mtx sync.Mutex
  23         src rand.Source
  24 }
  25
  26 // Next returns a new ID string. It is safe to call Next from multiple
  27 // goroutines.
  28 func (g *IDGenerator) Next() string {
  29         g.mtx.Lock()
  30         defer g.mtx.Unlock()
  31         if g.src == nil {
  32                 g.src = rand.NewSource(time.Now().UnixNano())
  33         }
  34         a, b := g.src.Int63(), g.src.Int63()
  35         id := strconv.FormatInt(a, 36) + strconv.FormatInt(b, 36)
  36         for len(id) > 20 {
  37                 id = id[:20]
  38         }
  39         return g.Prefix + id
  40 }
  41
  42 // AddRequestIDs wraps an http.Handler, adding an X-Request-Id header
  43 // to each request that doesn't already have one.
  44 func AddRequestIDs(h http.Handler) http.Handler {
  45         gen := &IDGenerator{Prefix: "req-"}
  46         return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
  47                 if req.Header.Get("X-Request-Id") == "" {
  48                         if req.Header == nil {
  49                                 req.Header = http.Header{}
  50                         }
  51                         req.Header.Set("X-Request-Id", gen.Next())
  52                 }
  53                 h.ServeHTTP(w, req)
  54         })
  55 }