Add memberMap to State to speed up member queries (#411)
* Fix #406: reconnect() can be called while still connected * Add memberMap to speed up member queries * Fix error return value and remove deletion * Fix GuildAdd member map initialization edge case
This commit is contained in:
parent
faa8ececc0
commit
42d1f62e58
1 changed files with 49 additions and 15 deletions
64
state.go
64
state.go
|
@ -42,6 +42,7 @@ type State struct {
|
||||||
|
|
||||||
guildMap map[string]*Guild
|
guildMap map[string]*Guild
|
||||||
channelMap map[string]*Channel
|
channelMap map[string]*Channel
|
||||||
|
memberMap map[string]map[string]*Member
|
||||||
}
|
}
|
||||||
|
|
||||||
// NewState creates an empty state.
|
// NewState creates an empty state.
|
||||||
|
@ -59,9 +60,18 @@ func NewState() *State {
|
||||||
TrackPresences: true,
|
TrackPresences: true,
|
||||||
guildMap: make(map[string]*Guild),
|
guildMap: make(map[string]*Guild),
|
||||||
channelMap: make(map[string]*Channel),
|
channelMap: make(map[string]*Channel),
|
||||||
|
memberMap: make(map[string]map[string]*Member),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (s *State) createMemberMap(guild *Guild) {
|
||||||
|
members := make(map[string]*Member)
|
||||||
|
for _, m := range guild.Members {
|
||||||
|
members[m.User.ID] = m
|
||||||
|
}
|
||||||
|
s.memberMap[guild.ID] = members
|
||||||
|
}
|
||||||
|
|
||||||
// GuildAdd adds a guild to the current world state, or
|
// GuildAdd adds a guild to the current world state, or
|
||||||
// updates it if it already exists.
|
// updates it if it already exists.
|
||||||
func (s *State) GuildAdd(guild *Guild) error {
|
func (s *State) GuildAdd(guild *Guild) error {
|
||||||
|
@ -77,6 +87,14 @@ func (s *State) GuildAdd(guild *Guild) error {
|
||||||
s.channelMap[c.ID] = c
|
s.channelMap[c.ID] = c
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// If this guild contains a new member slice, we must regenerate the member map so the pointers stay valid
|
||||||
|
if guild.Members != nil {
|
||||||
|
s.createMemberMap(guild)
|
||||||
|
} else if _, ok := s.memberMap[guild.ID]; !ok {
|
||||||
|
// Even if we have no new member slice, we still initialize the member map for this guild if it doesn't exist
|
||||||
|
s.memberMap[guild.ID] = make(map[string]*Member)
|
||||||
|
}
|
||||||
|
|
||||||
if g, ok := s.guildMap[guild.ID]; ok {
|
if g, ok := s.guildMap[guild.ID]; ok {
|
||||||
// We are about to replace `g` in the state with `guild`, but first we need to
|
// We are about to replace `g` in the state with `guild`, but first we need to
|
||||||
// make sure we preserve any fields that the `guild` doesn't contain from `g`.
|
// make sure we preserve any fields that the `guild` doesn't contain from `g`.
|
||||||
|
@ -271,14 +289,19 @@ func (s *State) MemberAdd(member *Member) error {
|
||||||
s.Lock()
|
s.Lock()
|
||||||
defer s.Unlock()
|
defer s.Unlock()
|
||||||
|
|
||||||
for i, m := range guild.Members {
|
members, ok := s.memberMap[member.GuildID]
|
||||||
if m.User.ID == member.User.ID {
|
if !ok {
|
||||||
guild.Members[i] = member
|
return ErrStateNotFound
|
||||||
return nil
|
}
|
||||||
}
|
|
||||||
|
m, ok := members[member.User.ID]
|
||||||
|
if !ok {
|
||||||
|
members[member.User.ID] = member
|
||||||
|
guild.Members = append(guild.Members, member)
|
||||||
|
} else {
|
||||||
|
*m = *member // Update the actual data, which will also update the member pointer in the slice
|
||||||
}
|
}
|
||||||
|
|
||||||
guild.Members = append(guild.Members, member)
|
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -296,6 +319,17 @@ func (s *State) MemberRemove(member *Member) error {
|
||||||
s.Lock()
|
s.Lock()
|
||||||
defer s.Unlock()
|
defer s.Unlock()
|
||||||
|
|
||||||
|
members, ok := s.memberMap[member.GuildID]
|
||||||
|
if !ok {
|
||||||
|
return ErrStateNotFound
|
||||||
|
}
|
||||||
|
|
||||||
|
_, ok = members[member.User.ID]
|
||||||
|
if !ok {
|
||||||
|
return ErrStateNotFound
|
||||||
|
}
|
||||||
|
delete(members, member.User.ID)
|
||||||
|
|
||||||
for i, m := range guild.Members {
|
for i, m := range guild.Members {
|
||||||
if m.User.ID == member.User.ID {
|
if m.User.ID == member.User.ID {
|
||||||
guild.Members = append(guild.Members[:i], guild.Members[i+1:]...)
|
guild.Members = append(guild.Members[:i], guild.Members[i+1:]...)
|
||||||
|
@ -312,18 +346,17 @@ func (s *State) Member(guildID, userID string) (*Member, error) {
|
||||||
return nil, ErrNilState
|
return nil, ErrNilState
|
||||||
}
|
}
|
||||||
|
|
||||||
guild, err := s.Guild(guildID)
|
|
||||||
if err != nil {
|
|
||||||
return nil, err
|
|
||||||
}
|
|
||||||
|
|
||||||
s.RLock()
|
s.RLock()
|
||||||
defer s.RUnlock()
|
defer s.RUnlock()
|
||||||
|
|
||||||
for _, m := range guild.Members {
|
members, ok := s.memberMap[guildID]
|
||||||
if m.User.ID == userID {
|
if !ok {
|
||||||
return m, nil
|
return nil, ErrStateNotFound
|
||||||
}
|
}
|
||||||
|
|
||||||
|
m, ok := members[userID]
|
||||||
|
if ok {
|
||||||
|
return m, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
return nil, ErrStateNotFound
|
return nil, ErrStateNotFound
|
||||||
|
@ -735,6 +768,7 @@ func (s *State) onReady(se *Session, r *Ready) (err error) {
|
||||||
|
|
||||||
for _, g := range s.Guilds {
|
for _, g := range s.Guilds {
|
||||||
s.guildMap[g.ID] = g
|
s.guildMap[g.ID] = g
|
||||||
|
s.createMemberMap(g)
|
||||||
|
|
||||||
for _, c := range g.Channels {
|
for _, c := range g.Channels {
|
||||||
s.channelMap[c.ID] = c
|
s.channelMap[c.ID] = c
|
||||||
|
|
Loading…
Reference in a new issue