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