From 316e8443312ff705e8215947cb34297d81c1bc34 Mon Sep 17 00:00:00 2001
From: Kyle Lemons <kyle.lemons@reddit.com>
Date: Fri, 26 Aug 2022 16:11:09 -0700
Subject: [PATCH] Allow changing the implementation of Interface at run-time
 (#562)

* Allow changing the implementation of Interface at run-time

This is primarily necessary during tests, as most production code won't need to change,
but this is based on an internal detail of atomic.Value, which asserts the underlying
type of the value being Store'd is always identical.  Since it doesn't see Interface,
it only sees interface{} (aka any), it actually cares about the concrete type.

There are various ways we can resolve this, covered in the benchmark,
and with the current compiler using atomic.Value with a concrete struct
type seems to be the winner.  This is faster than the older interface version
because the type assertion to a concrete type is faster than the type assertion
to an interface type.

* Be more cautious about the global load and add a test to ensure it doesn't panic
* Remove the 1.18-only (generic) benchmark since it's not actually _doing_ anything
* it's 1.19
---
 ecinterface/global.go      |  27 +++++++---
 ecinterface/global_test.go | 106 +++++++++++++++++++++++++++++++++++++
 2 files changed, 125 insertions(+), 8 deletions(-)
 create mode 100644 ecinterface/global_test.go

diff --git a/ecinterface/global.go b/ecinterface/global.go
index 89f3f42c4..7b15bebad 100644
--- a/ecinterface/global.go
+++ b/ecinterface/global.go
@@ -30,12 +30,23 @@ var Logger log.Wrapper
 // ErrGetBeforeSet is the error returned when Get is called before Set.
 var ErrGetBeforeSet = errors.New("ecinterface: Get called before Set is called")
 
-// actual type: Interface
+// current is the storage type of global.
+//
+// atomic.Value requires that the underlying concrete type remain constant.
+// If we try to store two different implementations of Interface, we will get a panic,
+// because Interface is promoted to any when you call Store.
+//
+// Thus, we use a `current{}` so that the concrete type is always the same.
+type current struct {
+	Interface
+}
+
+// actual type: current
 var global atomic.Value
 
 // Set sets the global edge context implementation.
 func Set(impl Interface) {
-	global.Store(impl)
+	global.Store(current{impl})
 }
 
 // Get returns the previously Set global edge context implementation.
@@ -49,13 +60,13 @@ func Set(impl Interface) {
 //
 // - Its ContextToHeader always return ("", false).
 func Get() Interface {
-	v := global.Load()
-	if impl, _ := v.(Interface); impl != nil {
-		return impl
+	stored := global.Load()
+	if stored == nil {
+		Logger.Log(context.Background(), ErrGetBeforeSet.Error())
+		getBeforeSet.Inc()
+		return nopImpl
 	}
-	Logger.Log(context.Background(), ErrGetBeforeSet.Error())
-	getBeforeSet.Inc()
-	return nopImpl
+	return stored.(current).Interface
 }
 
 type nop struct{}
diff --git a/ecinterface/global_test.go b/ecinterface/global_test.go
new file mode 100644
index 000000000..e726f3eb4
--- /dev/null
+++ b/ecinterface/global_test.go
@@ -0,0 +1,106 @@
+package ecinterface
+
+import (
+	"sync"
+	"sync/atomic"
+	"testing"
+)
+
+func TestGlobal_doesntPanic(t *testing.T) {
+	Get()
+}
+
+func BenchmarkAtomics(b *testing.B) {
+	// Example results:
+	//
+	// BenchmarkAtomics/rwmutex-10         	    85623068	    13.76 ns/op
+	// BenchmarkAtomics/atomicValue-10     	    283137508	     4.204 ns/op
+	// BenchmarkAtomics/atomicStructValue-10    703642080	     1.705 ns/op
+	// BenchmarkAtomics/atomicPointer-10        568728128	     2.070 ns/op
+	b.Run("current", func(b *testing.B) {
+		Set(nop{})
+		b.ReportAllocs()
+		b.ResetTimer()
+		var count int
+		for i := 0; i < b.N; i++ {
+			if Get() != nil {
+				count++
+			}
+		}
+		if count != b.N {
+			b.Fatalf("this is just to avoid eliding the call")
+		}
+	})
+	b.Run("rwmutex", func(b *testing.B) {
+		var scope struct {
+			rw sync.RWMutex
+			Interface
+		}
+		scope.Interface = nop{}
+		get := func() Interface {
+			scope.rw.RLock()
+			defer scope.rw.RUnlock()
+			return scope.Interface
+		}
+		b.ReportAllocs()
+		b.ResetTimer()
+		var count int
+		for i := 0; i < b.N; i++ {
+			if get() != nil {
+				count++
+			}
+		}
+		if count != b.N {
+			b.Fatalf("this is just to avoid eliding the call")
+		}
+	})
+	b.Run("atomicValue", func(b *testing.B) {
+		var global atomic.Value
+		global.Store(nop{})
+		b.ReportAllocs()
+		b.ResetTimer()
+		var count int
+		for i := 0; i < b.N; i++ {
+			if global.Load().(Interface) != nil {
+				count++
+			}
+		}
+		if count != b.N {
+			b.Fatalf("this is just to avoid eliding the call")
+		}
+	})
+	b.Run("atomicStructValue", func(b *testing.B) {
+		var global atomic.Value
+		type current struct{ Interface }
+		global.Store(current{nop{}})
+		b.ReportAllocs()
+		b.ResetTimer()
+		var count int
+		for i := 0; i < b.N; i++ {
+			if global.Load().(current).Interface != nil {
+				count++
+			}
+		}
+		if count != b.N {
+			b.Fatalf("this is just to avoid eliding the call")
+		}
+	})
+	/* If you have go1.19 with the new atomic APIs, you can throw this one in for fun:
+	b.Run("atomicPointer", func(b *testing.B) {
+		var global atomic.Pointer[Interface]
+		var impl Interface = nop{}
+		global.Store(&impl)
+		b.ReportAllocs()
+		b.ResetTimer()
+		var count int
+		for i := 0; i < b.N; i++ {
+			if *global.Load() != nil {
+				count++
+			}
+		}
+		if count != b.N {
+			b.Fatalf("this is just to avoid eliding the call")
+		}
+	})
+	*/
+}