Christian Schwarz
59 lines
1.2 KiB
Go
59 lines
1.2 KiB
Go
package rpc
|
|
|
|
import (
|
|
"context"
|
|
"encoding/base64"
|
|
"strings"
|
|
|
|
"github.com/google/uuid"
|
|
)
|
|
|
|
type requestIDGenerator struct{}
|
|
|
|
func newRequestIDGenerator() *requestIDGenerator {
|
|
return &requestIDGenerator{}
|
|
}
|
|
|
|
type RequestID struct{ s string }
|
|
|
|
func (r RequestID) String() string { return r.s }
|
|
|
|
type requestIDContextKey int
|
|
|
|
const (
|
|
requestIDContextKeyRequestID requestIDContextKey = 1 + iota
|
|
)
|
|
|
|
func (g *requestIDGenerator) newID() RequestID {
|
|
id := uuid.New()
|
|
var buf strings.Builder
|
|
enc := base64.NewEncoder(base64.RawStdEncoding, &buf)
|
|
n, err := enc.Write(id[:])
|
|
if err != nil {
|
|
panic(err)
|
|
} else if n != len(id) {
|
|
panic(n)
|
|
}
|
|
if err := enc.Close(); err != nil {
|
|
panic(err)
|
|
}
|
|
return RequestID{buf.String()}
|
|
}
|
|
|
|
func (g *requestIDGenerator) inject(ctx context.Context) context.Context {
|
|
return context.WithValue(ctx, requestIDContextKeyRequestID, g.newID())
|
|
}
|
|
|
|
func GetRequestID(ctx context.Context) (id RequestID, ok bool) {
|
|
id, ok = ctx.Value(requestIDContextKeyRequestID).(RequestID)
|
|
return id, ok
|
|
}
|
|
|
|
func MustGetRequestID(ctx context.Context) (id RequestID) {
|
|
id, ok := GetRequestID(ctx)
|
|
if !ok {
|
|
panic("calling context expectes request id to bet set")
|
|
}
|
|
return id
|
|
}
|