diff --git a/nsqlookupd/registration_db.go b/nsqlookupd/registration_db.go index d5a4825b3..ac74c7cc1 100644 --- a/nsqlookupd/registration_db.go +++ b/nsqlookupd/registration_db.go @@ -62,24 +62,50 @@ func NewRegistrationDB() *RegistrationDB { } // update high speed registrations map -// should call this func when registrationMap updated -func (r *RegistrationDB) fillHighSpeedRegistrations() { - r.highSpeedRegistrations = make(map[Registration]Registrations) - for k := range r.registrationMap { - key := Registration{k.Category, "*", ""} - if _, ok := r.highSpeedRegistrations[key]; !ok { - r.highSpeedRegistrations[key] = Registrations{} +// should call this func when registrationMap add Registration +func (r *RegistrationDB) addHighSpeedRegistration(k Registration) { + key := Registration{k.Category, "*", ""} + if _, ok := r.highSpeedRegistrations[key]; !ok { + r.highSpeedRegistrations[key] = Registrations{} + } + if k.IsMatch(key.Category, key.Key, key.SubKey) { + r.highSpeedRegistrations[key] = append(r.highSpeedRegistrations[key], k) + } + if k.SubKey != "" { + subKey := Registration{k.Category, k.Key, "*"} + if _, ok := r.highSpeedRegistrations[subKey]; !ok { + r.highSpeedRegistrations[subKey] = Registrations{} } - if k.IsMatch(key.Category, key.Key, key.SubKey) { - r.highSpeedRegistrations[key] = append(r.highSpeedRegistrations[key], k) + if k.IsMatch(subKey.Category, subKey.Key, subKey.SubKey) { + r.highSpeedRegistrations[subKey] = append(r.highSpeedRegistrations[subKey], k) } - if k.SubKey != "" { - subKey := Registration{k.Category, k.Key, "*"} - if _, ok := r.highSpeedRegistrations[subKey]; !ok { - r.highSpeedRegistrations[subKey] = Registrations{} + } +} + +// update high speed registrations map +// should call this func when registrationMap remove Registration +func (r *RegistrationDB) removeHighSpeedRegistration(k Registration) { + key := Registration{k.Category, "*", ""} + if registrations, ok := r.highSpeedRegistrations[key]; ok { + for i, registration := range registrations { + if registration == k { + r.highSpeedRegistrations[key] = append( + r.highSpeedRegistrations[key][:i], r.highSpeedRegistrations[key][i+1:]..., + ) + break } - if k.IsMatch(subKey.Category, subKey.Key, subKey.SubKey) { - r.highSpeedRegistrations[subKey] = append(r.highSpeedRegistrations[subKey], k) + } + } + if k.SubKey != "" { + subKey := Registration{k.Category, k.Key, "*"} + if registrations, ok := r.highSpeedRegistrations[subKey]; ok { + for i, registration := range registrations { + if registration == k { + r.highSpeedRegistrations[subKey] = append( + r.highSpeedRegistrations[subKey][:i], r.highSpeedRegistrations[subKey][i+1:]..., + ) + break + } } } } @@ -92,7 +118,7 @@ func (r *RegistrationDB) AddRegistration(k Registration) { _, ok := r.registrationMap[k] if !ok { r.registrationMap[k] = make(map[string]*Producer) - r.fillHighSpeedRegistrations() + r.addHighSpeedRegistration(k) } } @@ -103,7 +129,7 @@ func (r *RegistrationDB) AddProducer(k Registration, p *Producer) bool { _, ok := r.registrationMap[k] if !ok { r.registrationMap[k] = make(map[string]*Producer) - r.fillHighSpeedRegistrations() + r.addHighSpeedRegistration(k) } producers := r.registrationMap[k] _, found := producers[p.peerInfo.id] @@ -152,7 +178,7 @@ func (r *RegistrationDB) RemoveRegistration(k Registration) { r.Lock() defer r.Unlock() delete(r.registrationMap, k) - r.fillHighSpeedRegistrations() + r.removeHighSpeedRegistration(k) } func (r *RegistrationDB) needFilter(key string, subkey string) bool {