From bacfb41fdd06edf5294635d18ec387dee671a964 Mon Sep 17 00:00:00 2001
From: Peter Bourgon <peterbourgon@users.noreply.github.com>
Date: Thu, 4 Apr 2019 10:43:05 -0700
Subject: [PATCH] TestMonotonicSafe fails sometimes (#44)

* TestMonotonicSafe fails sometimes

* TestMonotonicSafe: vary seed and inc

* Check for MonotonicReader interface instead of *monotonic

* rm seed/inc variation in test, it's not relevant

* TestMonotonicSafe should use multiple goroutines

* t.Fatalf must be called in the same goroutine as the test
---
 ulid.go      | 27 +++++++++++++++++++--------
 ulid_test.go | 50 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 8 deletions(-)

diff --git a/ulid.go b/ulid.go
index fff09ac..c2f16ae 100644
--- a/ulid.go
+++ b/ulid.go
@@ -76,6 +76,15 @@ var (
 	ErrScanValue = errors.New("ulid: source value must be a string or byte slice")
 )
 
+// MonotonicReader is an interface that should yield monotonically increasing
+// entropy into the provided slice for all calls with the same ms parameter. If
+// a MonotonicReader is provided to the New constructor, its MonotonicRead
+// method will be used instead of Read.
+type MonotonicReader interface {
+	io.Reader
+	MonotonicRead(ms uint64, p []byte) error
+}
+
 // New returns an ULID with the given Unix milliseconds timestamp and an
 // optional entropy source. Use the Timestamp function to convert
 // a time.Time to Unix milliseconds.
@@ -93,7 +102,7 @@ func New(ms uint64, entropy io.Reader) (id ULID, err error) {
 	switch e := entropy.(type) {
 	case nil:
 		return id, err
-	case *monotonic:
+	case MonotonicReader:
 		err = e.MonotonicRead(ms, id[6:])
 	default:
 		_, err = io.ReadFull(e, id[6:])
@@ -487,9 +496,9 @@ func (id ULID) Value() (driver.Value, error) {
 // secure entropy bytes, then don't go under this default unless you know
 // what you're doing.
 //
-// The returned io.Reader isn't safe for concurrent use.
-func Monotonic(entropy io.Reader, inc uint64) io.Reader {
-	m := monotonic{
+// The returned type isn't safe for concurrent use.
+func Monotonic(entropy io.Reader, inc uint64) *MonotonicEntropy {
+	m := MonotonicEntropy{
 		Reader: bufio.NewReader(entropy),
 		inc:    inc,
 	}
@@ -505,7 +514,8 @@ func Monotonic(entropy io.Reader, inc uint64) io.Reader {
 	return &m
 }
 
-type monotonic struct {
+// MonotonicEntropy is an opaque type that provides monotonic entropy.
+type MonotonicEntropy struct {
 	io.Reader
 	ms      uint64
 	inc     uint64
@@ -514,7 +524,8 @@ type monotonic struct {
 	rng     *rand.Rand
 }
 
-func (m *monotonic) MonotonicRead(ms uint64, entropy []byte) (err error) {
+// MonotonicRead implements the MonotonicReader interface.
+func (m *MonotonicEntropy) MonotonicRead(ms uint64, entropy []byte) (err error) {
 	if !m.entropy.IsZero() && m.ms == ms {
 		err = m.increment()
 		m.entropy.AppendTo(entropy)
@@ -527,7 +538,7 @@ func (m *monotonic) MonotonicRead(ms uint64, entropy []byte) (err error) {
 
 // increment the previous entropy number with a random number
 // of up to m.inc (inclusive).
-func (m *monotonic) increment() error {
+func (m *MonotonicEntropy) increment() error {
 	if inc, err := m.random(); err != nil {
 		return err
 	} else if m.entropy.Add(inc) {
@@ -539,7 +550,7 @@ func (m *monotonic) increment() error {
 // random returns a uniform random value in [1, m.inc), reading entropy
 // from m.Reader. When m.inc == 0 || m.inc == 1, it returns 1.
 // Adapted from: https://golang.org/pkg/crypto/rand/#Int
-func (m *monotonic) random() (inc uint64, err error) {
+func (m *MonotonicEntropy) random() (inc uint64, err error) {
 	if m.inc <= 1 {
 		return 1, nil
 	}
diff --git a/ulid_test.go b/ulid_test.go
index 048adbf..b6f930a 100644
--- a/ulid_test.go
+++ b/ulid_test.go
@@ -21,6 +21,7 @@ import (
 	"math"
 	"math/rand"
 	"strings"
+	"sync"
 	"testing"
 	"testing/iotest"
 	"testing/quick"
@@ -581,6 +582,55 @@ func TestMonotonicOverflow(t *testing.T) {
 	}
 }
 
+func TestMonotonicSafe(t *testing.T) {
+	t.Parallel()
+
+	var (
+		src       = rand.NewSource(time.Now().UnixNano())
+		entropy   = rand.New(src)
+		monotonic = ulid.Monotonic(entropy, 0)
+		safe      = &safeMonotonicReader{MonotonicReader: monotonic}
+		t0        = ulid.Timestamp(time.Now())
+	)
+
+	errs := make(chan error, 100)
+	for i := 0; i < cap(errs); i++ {
+		go func() {
+			u0 := ulid.MustNew(t0, safe)
+			u1 := ulid.MustNew(t0, safe)
+			for j := 0; j < 1024; j++ {
+				u0, u1 = u1, ulid.MustNew(t0, safe)
+				if u0.String() >= u1.String() {
+					errs <- fmt.Errorf(
+						"%s (%d %x) >= %s (%d %x)",
+						u0.String(), u0.Time(), u0.Entropy(),
+						u1.String(), u1.Time(), u1.Entropy(),
+					)
+					return
+				}
+			}
+			errs <- nil
+		}()
+	}
+	for i := 0; i < cap(errs); i++ {
+		if err := <-errs; err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+type safeMonotonicReader struct {
+	mtx sync.Mutex
+	ulid.MonotonicReader
+}
+
+func (r *safeMonotonicReader) MonotonicRead(ms uint64, p []byte) (err error) {
+	r.mtx.Lock()
+	err = r.MonotonicReader.MonotonicRead(ms, p)
+	r.mtx.Unlock()
+	return err
+}
+
 func BenchmarkNew(b *testing.B) {
 	benchmarkMakeULID(b, func(timestamp uint64, entropy io.Reader) {
 		_, _ = ulid.New(timestamp, entropy)