Moreatomic: Replaced invalid mutex with modified sync.Map impl

gateway/ready.go

@@ -200,7 +200,7 @@ type (
 	// ReadySupplementalEvent is the struct for a READY_SUPPLEMENTAL event,
 	// which is an undocumented event.
 	ReadySupplementalEvent struct {
-		Guilds          []discord.Guild        `json:"guilds"` // only have ID and VoiceStates
+		Guilds          []GuildCreateEvent     `json:"guilds"` // only have ID and VoiceStates
 		MergedMembers   [][]SupplementalMember `json:"merged_members"`
 		MergedPresences MergedPresences        `json:"merged_presences"`
 	}

internal/moreatomic/syncmap.go

@@ -1,61 +1,39 @@
 package moreatomic
 
 import (
-	"sync"
 	"sync/atomic"
+
+	"github.com/diamondburned/arikawa/v2/internal/moreatomic/syncmod"
 )
 
 // Map is a thread-safe map that is a wrapper around sync.Map with slight API
 // additions.
 type Map struct {
-	smap atomic.Value
+	val  atomic.Value
 	ctor func() interface{}
 }
 
-type sentinelType struct{}
-
-var sentinel = sentinelType{}
-
 func NewMap(ctor func() interface{}) *Map {
-	smap := atomic.Value{}
-	smap.Store(&sync.Map{})
-	return &Map{smap, ctor}
+	sm := &Map{ctor: ctor}
+	sm.Reset()
+	return sm
 }
 
 // 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.val.Store(&syncmod.Map{New: sm.ctor})
 	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)
-
-	lv, loaded = smap.LoadOrStore(k, sentinel)
-	if !loaded {
-		lv = sm.ctor()
-		smap.Store(k, lv)
-	}
-
-	return
+	return sm.val.Load().(*syncmod.Map).LoadOrStore(k)
 }
 
 // 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)
-
-	for {
-		lv, ok = smap.Load(k)
-		if !ok {
-			return nil, false
-		}
-
-		if lv != sentinel {
-			return lv, true
-		}
-	}
+	return sm.val.Load().(*syncmod.Map).Load(k)
 }

internal/moreatomic/syncmod/syncmod.go

@@ -0,0 +1,263 @@
+// Copyright 2016 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package syncmod contains a clone of package sync's map.go file with unused
+// methods removed and some tweaks with LoadOrStore.
+package syncmod
+
+import (
+	"sync"
+	"sync/atomic"
+	"unsafe"
+)
+
+// Map is like a Go map[interface{}]interface{} but is safe for concurrent use
+// by multiple goroutines without additional locking or coordination.
+// Loads, stores, and deletes run in amortized constant time.
+//
+// The Map type is specialized. Most code should use a plain Go map instead,
+// with separate locking or coordination, for better type safety and to make it
+// easier to maintain other invariants along with the map content.
+//
+// The Map type is optimized for two common use cases: (1) when the entry for a given
+// key is only ever written once but read many times, as in caches that only grow,
+// or (2) when multiple goroutines read, write, and overwrite entries for disjoint
+// sets of keys. In these two cases, use of a Map may significantly reduce lock
+// contention compared to a Go map paired with a separate Mutex or RWMutex.
+//
+// The zero Map is empty and ready for use. A Map must not be copied after first use.
+type Map struct {
+	New func() interface{}
+
+	mu sync.Mutex
+
+	// read contains the portion of the map's contents that are safe for
+	// concurrent access (with or without mu held).
+	//
+	// The read field itself is always safe to load, but must only be stored with
+	// mu held.
+	//
+	// Entries stored in read may be updated concurrently without mu, but updating
+	// a previously-expunged entry requires that the entry be copied to the dirty
+	// map and unexpunged with mu held.
+	read atomic.Value // readOnly
+
+	// dirty contains the portion of the map's contents that require mu to be
+	// held. To ensure that the dirty map can be promoted to the read map quickly,
+	// it also includes all of the non-expunged entries in the read map.
+	//
+	// Expunged entries are not stored in the dirty map. An expunged entry in the
+	// clean map must be unexpunged and added to the dirty map before a new value
+	// can be stored to it.
+	//
+	// If the dirty map is nil, the next write to the map will initialize it by
+	// making a shallow copy of the clean map, omitting stale entries.
+	dirty map[interface{}]*entry
+
+	// misses counts the number of loads since the read map was last updated that
+	// needed to lock mu to determine whether the key was present.
+	//
+	// Once enough misses have occurred to cover the cost of copying the dirty
+	// map, the dirty map will be promoted to the read map (in the unamended
+	// state) and the next store to the map will make a new dirty copy.
+	misses int
+}
+
+// readOnly is an immutable struct stored atomically in the Map.read field.
+type readOnly struct {
+	m       map[interface{}]*entry
+	amended bool // true if the dirty map contains some key not in m.
+}
+
+// expunged is an arbitrary pointer that marks entries which have been deleted
+// from the dirty map.
+var expunged = unsafe.Pointer(new(interface{}))
+
+// An entry is a slot in the map corresponding to a particular key.
+type entry struct {
+	// p points to the interface{} value stored for the entry.
+	//
+	// If p == nil, the entry has been deleted and m.dirty == nil.
+	//
+	// If p == expunged, the entry has been deleted, m.dirty != nil, and the entry
+	// is missing from m.dirty.
+	//
+	// Otherwise, the entry is valid and recorded in m.read.m[key] and, if m.dirty
+	// != nil, in m.dirty[key].
+	//
+	// An entry can be deleted by atomic replacement with nil: when m.dirty is
+	// next created, it will atomically replace nil with expunged and leave
+	// m.dirty[key] unset.
+	//
+	// An entry's associated value can be updated by atomic replacement, provided
+	// p != expunged. If p == expunged, an entry's associated value can be updated
+	// only after first setting m.dirty[key] = e so that lookups using the dirty
+	// map find the entry.
+	p unsafe.Pointer // *interface{}
+}
+
+func newEntry(i interface{}) *entry {
+	return &entry{p: unsafe.Pointer(&i)}
+}
+
+// Load returns the value stored in the map for a key, or nil if no
+// value is present.
+// The ok result indicates whether value was found in the map.
+func (m *Map) Load(key interface{}) (value interface{}, ok bool) {
+	read, _ := m.read.Load().(readOnly)
+	e, ok := read.m[key]
+	if !ok && read.amended {
+		m.mu.Lock()
+		// Avoid reporting a spurious miss if m.dirty got promoted while we were
+		// blocked on m.mu. (If further loads of the same key will not miss, it's
+		// not worth copying the dirty map for this key.)
+		read, _ = m.read.Load().(readOnly)
+		e, ok = read.m[key]
+		if !ok && read.amended {
+			e, ok = m.dirty[key]
+			// Regardless of whether the entry was present, record a miss: this key
+			// will take the slow path until the dirty map is promoted to the read
+			// map.
+			m.missLocked()
+		}
+		m.mu.Unlock()
+	}
+	if !ok {
+		return nil, false
+	}
+	return e.load()
+}
+
+func (e *entry) load() (value interface{}, ok bool) {
+	p := atomic.LoadPointer(&e.p)
+	if p == nil || p == expunged {
+		return nil, false
+	}
+	return *(*interface{})(p), true
+}
+
+// unexpungeLocked ensures that the entry is not marked as expunged.
+//
+// If the entry was previously expunged, it must be added to the dirty map
+// before m.mu is unlocked.
+func (e *entry) unexpungeLocked() (wasExpunged bool) {
+	return atomic.CompareAndSwapPointer(&e.p, expunged, nil)
+}
+
+// LoadOrStore returns the existing value for the key if present.
+// Otherwise, it stores and returns the given value.
+// The loaded result is true if the value was loaded, false if stored.
+func (m *Map) LoadOrStore(k interface{}) (actual interface{}, loaded bool) {
+	// Avoid locking if it's a clean hit.
+	read, _ := m.read.Load().(readOnly)
+	if e, ok := read.m[k]; ok {
+		actual, loaded, ok = e.tryLoadOrStore(nil, m.New)
+		if ok {
+			return actual, loaded
+		}
+	}
+
+	m.mu.Lock()
+	read, _ = m.read.Load().(readOnly)
+	if e, ok := read.m[k]; ok {
+		if e.unexpungeLocked() {
+			m.dirty[k] = e
+		}
+		actual, loaded, _ = e.tryLoadOrStore(actual, m.New)
+	} else if e, ok := m.dirty[k]; ok {
+		actual, loaded, _ = e.tryLoadOrStore(actual, m.New)
+		m.missLocked()
+	} else {
+		if !read.amended {
+			// We're adding the first new key to the dirty map.
+			// Make sure it is allocated and mark the read-only map as incomplete.
+			m.dirtyLocked()
+			m.read.Store(readOnly{m: read.m, amended: true})
+		}
+		// This will likely allocate if the first tryLoadOrStore sees an
+		// expunged value and this else branch is hit.
+		if actual == nil {
+			actual = m.New()
+		}
+		m.dirty[k] = newEntry(actual)
+		loaded = false
+	}
+	m.mu.Unlock()
+
+	return actual, loaded
+}
+
+// tryLoadOrStore atomically loads or stores a value if the entry is not
+// expunged.
+//
+// If the entry is expunged, tryLoadOrStore leaves the entry unchanged and
+// returns with ok==false.
+func (e *entry) tryLoadOrStore(
+	i interface{}, newFn func() interface{}) (actual interface{}, loaded, ok bool) {
+
+	p := atomic.LoadPointer(&e.p)
+	if p == expunged {
+		return nil, false, false
+	}
+	if p != nil {
+		return *(*interface{})(p), true, true
+	}
+
+	if i == nil {
+		i = newFn()
+	}
+
+	// Copy the interface after the first load to make this method more amenable
+	// to escape analysis: if we hit the "load" path or the entry is expunged, we
+	// shouldn't bother heap-allocating.
+	ic := i
+
+	for {
+		if atomic.CompareAndSwapPointer(&e.p, nil, unsafe.Pointer(&ic)) {
+			return i, false, true
+		}
+		p = atomic.LoadPointer(&e.p)
+		if p == expunged {
+			return i, false, false
+		}
+		if p != nil {
+			return *(*interface{})(p), true, true
+		}
+	}
+}
+
+func (m *Map) missLocked() {
+	m.misses++
+	if m.misses < len(m.dirty) {
+		return
+	}
+	m.read.Store(readOnly{m: m.dirty})
+	m.dirty = nil
+	m.misses = 0
+}
+
+func (m *Map) dirtyLocked() {
+	if m.dirty != nil {
+		return
+	}
+
+	read, _ := m.read.Load().(readOnly)
+	m.dirty = make(map[interface{}]*entry, len(read.m))
+	for k, e := range read.m {
+		if !e.tryExpungeLocked() {
+			m.dirty[k] = e
+		}
+	}
+}
+
+func (e *entry) tryExpungeLocked() (isExpunged bool) {
+	p := atomic.LoadPointer(&e.p)
+	for p == nil {
+		if atomic.CompareAndSwapPointer(&e.p, nil, expunged) {
+			return true
+		}
+		p = atomic.LoadPointer(&e.p)
+	}
+	return p == expunged
+}

state/state_events.go

@@ -82,7 +82,43 @@ func (s *State) onEvent(iface interface{}) {
 		s.readyMu.Unlock()
 
 	case *gateway.ReadySupplementalEvent:
-		// TODO
+		// Handle guilds
+		for _, guild := range ev.Guilds {
+			// Handle guild voice states
+			for _, v := range guild.VoiceStates {
+				if err := s.Cabinet.VoiceStateSet(guild.ID, v); err != nil {
+					s.stateErr(err, "failed to set guild voice state in Ready Supplemental")
+				}
+			}
+		}
+
+		for _, friend := range ev.MergedPresences.Friends {
+			sPresence := gateway.ConvertSupplementalPresence(friend)
+			if err := s.Cabinet.PresenceSet(0, sPresence); err != nil {
+				s.stateErr(err, "failed to set friend presence in Ready Supplemental")
+			}
+		}
+
+		// Discord uses weird indexing, so we'll need the Guilds slice.
+		ready := s.Ready()
+
+		for i := 0; i < len(ready.Guilds) && i < len(ev.MergedMembers); i++ {
+			guild := ready.Guilds[i]
+
+			for _, member := range ev.MergedMembers[i] {
+				sMember := gateway.ConvertSupplementalMember(member)
+				if err := s.Cabinet.MemberSet(guild.ID, sMember); err != nil {
+					s.stateErr(err, "failed to set friend presence in Ready Supplemental")
+				}
+			}
+
+			for _, member := range ev.MergedPresences.Guilds[i] {
+				sPresence := gateway.ConvertSupplementalPresence(member)
+				if err := s.Cabinet.PresenceSet(guild.ID, sPresence); err != nil {
+					s.stateErr(err, "failed to set member presence in Ready Supplemental")
+				}
+			}
+		}
 
 	case *gateway.GuildCreateEvent:
 		s.batchLog(storeGuildCreate(s.Cabinet, ev))

state/store/defaultstore/channel.go

@@ -13,7 +13,9 @@ type Channel struct {
 
 	// Channel references must be protected under the same mutex.
 
-	privates map[discord.UserID]*discord.Channel
+	privates   map[discord.UserID]*discord.Channel
+	privateChs []*discord.Channel
+
 	channels map[discord.ChannelID]*discord.Channel
 	guildChs map[discord.GuildID][]*discord.Channel
 }
@@ -90,21 +92,19 @@ func (s *Channel) PrivateChannels() ([]discord.Channel, error) {
 	s.mut.RLock()
 	defer s.mut.RUnlock()
 
-	if len(s.privates) == 0 {
+	if len(s.privateChs) == 0 {
 		return nil, store.ErrNotFound
 	}
 
-	var channels = make([]discord.Channel, 0, len(s.privates))
-	for _, ch := range s.privates {
-		channels = append(channels, *ch)
+	var channels = make([]discord.Channel, len(s.privateChs))
+	for i, ch := range s.privateChs {
+		channels[i] = *ch
 	}
 
 	return channels, nil
 }
 
-// ChannelSet sets the Direct Message or Guild channl into the state. If the
-// channel doesn't have 1 (one) DMRecipients, then it must have a valid GuildID,
-// otherwise an error will be returned.
+// ChannelSet sets the Direct Message or Guild channel into the state.
 func (s *Channel) ChannelSet(channel discord.Channel) error {
 	s.mut.Lock()
 	defer s.mut.Unlock()
@@ -115,20 +115,26 @@ func (s *Channel) ChannelSet(channel discord.Channel) error {
 		return nil
 	}
 
-	if len(channel.DMRecipients) == 1 {
+	switch channel.Type {
+	case discord.DirectMessage:
+		// Safety bound check.
+		if len(channel.DMRecipients) != 1 {
+			return errors.New("DirectMessage channel does not have 1 recipient")
+		}
 		s.privates[channel.DMRecipients[0].ID] = &channel
+		fallthrough
+	case discord.GroupDM:
+		s.privateChs = append(s.privateChs, &channel)
 		s.channels[channel.ID] = &channel
 		return nil
 	}
 
-	// Invalid channel case, as we need the GuildID to search for this channel.
+	// Ensure that if the channel is not a DM or group DM channel, then it must
+	// have a valid guild ID.
 	if !channel.GuildID.IsValid() {
 		return errors.New("invalid guildID for guild channel")
 	}
 
-	// Always ensure that if the channel is in the slice, then it will be in the
-	// map.
-
 	s.channels[channel.ID] = &channel
 
 	channels, _ := s.guildChs[channel.GuildID]
@@ -142,38 +148,57 @@ func (s *Channel) ChannelRemove(channel discord.Channel) error {
 	s.mut.Lock()
 	defer s.mut.Unlock()
 
+	// Wipe the channel off the channel ID index.
 	delete(s.channels, channel.ID)
 
-	if len(channel.DMRecipients) == 1 {
+	// Wipe the channel off the DM recipient index, if available.
+	switch channel.Type {
+	case discord.DirectMessage:
+		// Safety bound check.
+		if len(channel.DMRecipients) != 1 {
+			return errors.New("DirectMessage channel does not have 1 recipient")
+		}
 		delete(s.privates, channel.DMRecipients[0].ID)
+		fallthrough
+	case discord.GroupDM:
+		for i, priv := range s.privateChs {
+			if priv.ID == channel.ID {
+				s.privateChs = removeChannel(s.privateChs, i)
+				break
+			}
+		}
 		return nil
 	}
 
+	// Wipe the channel off the guilds index, if available.
 	channels, ok := s.guildChs[channel.GuildID]
 	if !ok {
 		return nil
 	}
 
 	for i, ch := range channels {
-		if ch.ID != channel.ID {
-			continue
+		if ch.ID == channel.ID {
+			s.guildChs[channel.GuildID] = removeChannel(channels, i)
+			break
 		}
-
-		// Fast unordered delete. Not sure if there's a benefit in doing
-		// this over using a map, but I guess the memory usage is less and
-		// there's no copying.
-
-		// Move the last channel to the current channel, set the last
-		// channel there to a nil value to unreference its children, then
-		// slice the last channel off.
-		channels[i] = channels[len(channels)-1]
-		channels[len(channels)-1] = nil
-		channels = channels[:len(channels)-1]
-
-		s.guildChs[channel.GuildID] = channels
-
-		break
 	}
 
 	return nil
 }
+
+// removeChannel removes the given channel with the index from the given
+// channels slice in an unordered fashion.
+func removeChannel(channels []*discord.Channel, i int) []*discord.Channel {
+	// Fast unordered delete. Not sure if there's a benefit in doing
+	// this over using a map, but I guess the memory usage is less and
+	// there's no copying.
+
+	// Move the last channel to the current channel, set the last
+	// channel there to a nil value to unreference its children, then
+	// slice the last channel off.
+	channels[i] = channels[len(channels)-1]
+	channels[len(channels)-1] = nil
+	channels = channels[:len(channels)-1]
+
+	return channels
+}