update highspeed map during each mutation

nsqio · Jan 1, 2021 · d134659 · d134659
1 parent a58cae7
commit d134659
Showing 1 changed file with 44 additions and 18 deletions.
diff --git 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 {