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:
Erik McClure 2017-07-29 08:09:06 -07:00 committed by Chris Rhodes
parent faa8ececc0
commit 42d1f62e58

View file

@ -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) guild.Members = append(guild.Members, member)
} else {
*m = *member // Update the actual data, which will also update the member pointer in the slice
}
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