github.com/qiniu/x@v1.11.9/reqid/reqid.go (about)

     1  package reqid
     2  
     3  import (
     4  	"encoding/base64"
     5  	"encoding/binary"
     6  	"net/http"
     7  	"time"
     8  
     9  	. "context"
    10  )
    11  
    12  // --------------------------------------------------------------------
    13  
    14  var pid = uint32(time.Now().UnixNano() % 4294967291)
    15  
    16  func genReqId() string {
    17  	var b [12]byte
    18  	binary.LittleEndian.PutUint32(b[:], pid)
    19  	binary.LittleEndian.PutUint64(b[4:], uint64(time.Now().UnixNano()))
    20  	return base64.URLEncoding.EncodeToString(b[:])
    21  }
    22  
    23  // --------------------------------------------------------------------
    24  
    25  type key int // key is unexported and used for Context
    26  
    27  const (
    28  	reqidKey key = 0
    29  )
    30  
    31  func NewContext(ctx Context, reqid string) Context {
    32  	return WithValue(ctx, reqidKey, reqid)
    33  }
    34  
    35  func NewContextWith(ctx Context, w http.ResponseWriter, req *http.Request) Context {
    36  	reqid := req.Header.Get("X-Reqid")
    37  	if reqid == "" {
    38  		reqid = genReqId()
    39  		req.Header.Set("X-Reqid", reqid)
    40  	}
    41  	h := w.Header()
    42  	h.Set("X-Reqid", reqid)
    43  	return WithValue(ctx, reqidKey, reqid)
    44  }
    45  
    46  func FromContext(ctx Context) (reqid string, ok bool) {
    47  	reqid, ok = ctx.Value(reqidKey).(string)
    48  	return
    49  }
    50  
    51  // --------------------------------------------------------------------