From 3ddb472644ba2a4b1b982ba58006117ac63f330f Mon Sep 17 00:00:00 2001
From: diamondburned <datutbrus@gmail.com>
Date: Fri, 18 Dec 2020 23:47:47 -0800
Subject: [PATCH] Moreatomic: Cave in and use upgradable mutex over sync.Map

---
 go.mod                         |  1 +
 go.sum                         |  2 ++
 internal/moreatomic/syncmap.go | 43 ++++++++++++++++------------------
 3 files changed, 23 insertions(+), 23 deletions(-)

diff --git a/go.mod b/go.mod
index 6d8e5e6..d1656fd 100644
--- a/go.mod
+++ b/go.mod
@@ -5,6 +5,7 @@ go 1.13
 require (
 	github.com/gorilla/schema v1.2.0
 	github.com/gorilla/websocket v1.4.2
+	github.com/kawasin73/umutex v0.2.1 // indirect
 	github.com/pkg/errors v0.9.1
 	golang.org/x/crypto v0.0.0-20201016220609-9e8e0b390897
 	golang.org/x/sys v0.0.0-20201018230417-eeed37f84f13 // indirect
diff --git a/go.sum b/go.sum
index 07e39d5..e55a9f1 100644
--- a/go.sum
+++ b/go.sum
@@ -2,6 +2,8 @@ github.com/gorilla/schema v1.2.0 h1:YufUaxZYCKGFuAq3c96BOhjgd5nmXiOY9NGzF247Tsc=
 github.com/gorilla/schema v1.2.0/go.mod h1:kgLaKoK1FELgZqMAVxx/5cbj0kT+57qxUrAlIO2eleU=
 github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0Ufc=
 github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
+github.com/kawasin73/umutex v0.2.1 h1:Onkzz3LKs1HThskVwdhhBocqdRQqwCZ03quDJzuPzPo=
+github.com/kawasin73/umutex v0.2.1/go.mod h1:A02N2muKVFMvFlp5c+hBycgdH964YtieGs+7mYB16NU=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2 h1:VklqNMn3ovrHsnt90PveolxSbWFaJdECFbxSq0Mqo2M=
diff --git a/internal/moreatomic/syncmap.go b/internal/moreatomic/syncmap.go
index 6a68cdb..ec22d7c 100644
--- a/internal/moreatomic/syncmap.go
+++ b/internal/moreatomic/syncmap.go
@@ -1,14 +1,12 @@
 package moreatomic
 
-import (
-	"sync"
-	"sync/atomic"
-)
+import "github.com/kawasin73/umutex"
 
 // Map is a thread-safe map that is a wrapper around sync.Map with slight API
 // additions.
 type Map struct {
-	smap atomic.Value
+	upmu umutex.UMutex
+	smap map[interface{}]interface{}
 	ctor func() interface{}
 }
 
@@ -17,27 +15,33 @@ type sentinelType struct{}
 var sentinel = sentinelType{}
 
 func NewMap(ctor func() interface{}) *Map {
-	smap := atomic.Value{}
-	smap.Store(&sync.Map{})
-	return &Map{smap, ctor}
+	return &Map{
+		smap: map[interface{}]interface{}{},
+		ctor: ctor,
+	}
 }
 
 // Reset swaps the internal map out with a fresh one, dropping the old map. This
 // method never errors.
 func (sm *Map) Reset() error {
-	sm.smap.Store(&sync.Map{})
+	sm.upmu.Lock()
+	sm.smap = map[interface{}]interface{}{}
+	sm.upmu.Unlock()
 	return nil
 }
 
 // LoadOrStore loads an existing value or stores a new value created from the
 // given constructor then return that value.
 func (sm *Map) LoadOrStore(k interface{}) (lv interface{}, loaded bool) {
-	smap := sm.smap.Load().(*sync.Map)
+	sm.upmu.RLock()
+	defer sm.upmu.RUnlock()
 
-	lv, loaded = smap.LoadOrStore(k, sentinel)
+	lv, loaded = sm.smap[k]
 	if !loaded {
 		lv = sm.ctor()
-		smap.Store(k, lv)
+
+		sm.upmu.Upgrade()
+		sm.smap[k] = lv
 	}
 
 	return
@@ -46,16 +50,9 @@ func (sm *Map) LoadOrStore(k interface{}) (lv interface{}, loaded bool) {
 // Load loads an existing value; it returns ok set to false if there is no
 // value with that key.
 func (sm *Map) Load(k interface{}) (lv interface{}, ok bool) {
-	smap := sm.smap.Load().(*sync.Map)
+	sm.upmu.RLock()
+	defer sm.upmu.RUnlock()
 
-	for {
-		lv, ok = smap.Load(k)
-		if !ok {
-			return nil, false
-		}
-
-		if lv != sentinel {
-			return lv, true
-		}
-	}
+	lv, ok = sm.smap[k]
+	return
 }