From 35abdefe61c21d00a95d8ee40311ebe74e2d9e00 Mon Sep 17 00:00:00 2001
From: Ryan Lang <rlang@me.com>
Date: Sun, 1 Mar 2020 19:05:58 -0800
Subject: [PATCH] new xray id generator

---
 xray.go      | 67 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 xray_test.go | 15 ++++++++++++
 2 files changed, 82 insertions(+)

diff --git a/xray.go b/xray.go
index caba398..47d2b96 100644
--- a/xray.go
+++ b/xray.go
@@ -16,14 +16,18 @@ package aws
 
 import (
 	"context"
+	crand "crypto/rand"
+	"encoding/binary"
 	"fmt"
 	"io"
 	"io/ioutil"
 	"log"
+	"math/rand"
 	"net/http"
 	"os"
 	"regexp"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -465,3 +469,66 @@ func (e *Exporter) makeSegment(span *trace.SpanData) segment {
 
 	return s
 }
+
+type xrayIDGenerator struct {
+	sync.Mutex
+
+	// Please keep these as the first fields
+	// so that these 8 byte fields will be aligned on addresses
+	// divisible by 8, on both 32-bit and 64-bit machines when
+	// performing atomic increments and accesses.
+	// See:
+	// * https://github.com/census-instrumentation/opencensus-go/issues/587
+	// * https://github.com/census-instrumentation/opencensus-go/issues/865
+	// * https://golang.org/pkg/sync/atomic/#pkg-note-BUG
+	nextSpanID uint64
+	spanIDInc  uint64
+
+	traceIDAdd  [2]uint64
+	traceIDRand *rand.Rand
+}
+
+func NewIDGenerator() *xrayIDGenerator {
+	gen := &xrayIDGenerator{}
+	// initialize traceID and spanID generators.
+
+	var rngSeed int64
+	for _, p := range []interface{}{
+		&rngSeed, &gen.traceIDAdd, &gen.nextSpanID, &gen.spanIDInc,
+	} {
+		binary.Read(crand.Reader, binary.LittleEndian, p)
+	}
+	gen.traceIDRand = rand.New(rand.NewSource(rngSeed))
+	gen.spanIDInc |= 1
+
+	return gen
+
+}
+
+// NewSpanID returns a non-zero span ID from a randomly-chosen sequence.
+func (gen *xrayIDGenerator) NewSpanID() [8]byte {
+	var id uint64
+	for id == 0 {
+		id = atomic.AddUint64(&gen.nextSpanID, gen.spanIDInc)
+	}
+	var sid [8]byte
+	binary.LittleEndian.PutUint64(sid[:], id)
+	return sid
+}
+
+// NewTraceID returns a non-zero trace ID from a randomly-chosen sequence.
+// mu should be held while this function is called. convertToAmazonTraceID
+// expects first 4 bytes to be current unix timestamp
+func (gen *xrayIDGenerator) NewTraceID() [16]byte {
+	var tid [16]byte
+	// Construct the trace ID from two outputs of traceIDRand, with a constant
+	// added to each half for additional entropy.
+	gen.Lock()
+
+	binary.BigEndian.PutUint32(tid[0:4], uint32(time.Now().Unix()))
+	binary.BigEndian.PutUint32(tid[4:8], gen.traceIDRand.Uint32())
+	binary.BigEndian.PutUint64(tid[8:16], gen.traceIDRand.Uint64()+gen.traceIDAdd[1])
+
+	gen.Unlock()
+	return tid
+}
diff --git a/xray_test.go b/xray_test.go
index a8787e9..3e52b5a 100644
--- a/xray_test.go
+++ b/xray_test.go
@@ -565,3 +565,18 @@ func TestLookupRegionFromMetaData(t *testing.T) {
 		})
 	}
 }
+
+func TestIDGenerator(t *testing.T) {
+	gen := NewIDGenerator()
+
+	sid := gen.NewSpanID()
+	tid := gen.NewTraceID()
+
+	if len(sid) != 8 {
+		t.Errorf("span id wrong length; want 8, got %d", len(sid))
+	}
+
+	if len(tid) != 16 {
+		t.Errorf("span id wrong length; want 16, got %d", len(tid))
+	}
+}