Ensure FLUSHDB clears all hook memory references

Issue #685
tidwall · May 11, 2023 · 5642fc4 · 5642fc4
1 parent a08c55b
commit 5642fc4
Show file tree

Hide file tree

Showing 2 changed files with 77 additions and 56 deletions.
diff --git a/internal/server/crud.go b/internal/server/crud.go
@@ -412,6 +412,15 @@ func (s *Server) cmdPDEL(msg *Message) (resp.Value, commandDetails, error) {
 	return res, d, nil
 }
 
+func (s *Server) cmdDROPop(key string) *collection.Collection {
+	col, _ := s.cols.Get(key)
+	if col != nil {
+		s.cols.Delete(key)
+	}
+	s.groupDisconnectCollection(key)
+	return col
+}
+
 // DROP key
 func (s *Server) cmdDROP(msg *Message) (resp.Value, commandDetails, error) {
 	start := time.Now()
@@ -425,12 +434,7 @@ func (s *Server) cmdDROP(msg *Message) (resp.Value, commandDetails, error) {
 	key := args[1]
 
 	// >> Operation
-
-	col, _ := s.cols.Get(key)
-	if col != nil {
-		s.cols.Delete(key)
-	}
-	s.groupDisconnectCollection(key)
+	col := s.cmdDROPop(key)
 
 	// >> Response
 
@@ -542,6 +546,39 @@ func (s *Server) cmdFLUSHDB(msg *Message) (resp.Value, commandDetails, error) {
 	// >> Operation
 
 	// clear the entire database
+
+	// drop each collection
+	keys := s.cols.Keys()
+	for _, key := range keys {
+		s.cmdDROPop(key)
+	}
+
+	// delete all channels
+	var names []string
+	s.hooks.Ascend(nil, func(item any) bool {
+		hook := item.(*Hook)
+		if hook.channel {
+			names = append(names, hook.Name)
+		}
+		return true
+	})
+	for _, name := range names {
+		s.cmdDELHOOKop(name, true)
+	}
+
+	// delete all hooks
+	names = names[:0]
+	s.hooks.Ascend(nil, func(item any) bool {
+		hook := item.(*Hook)
+		if !hook.channel {
+			names = append(names, hook.Name)
+		}
+		return true
+	})
+	for _, name := range names {
+		s.cmdDELHOOKop(name, false)
+	}
+
 	s.cols.Clear()
 	s.groupHooks.Clear()
 	s.groupObjects.Clear()

diff --git a/internal/server/hooks.go b/internal/server/hooks.go
@@ -240,6 +240,37 @@ func byHookExpires(a, b interface{}) bool {
 	return ha.Name < hb.Name
 }
 
+func (s *Server) cmdDELHOOKop(name string, channel bool) (updated bool) {
+	hook, _ := s.hooks.Get(&Hook{Name: name}).(*Hook)
+	if hook == nil || hook.channel != channel {
+		return false
+	}
+	hook.Close()
+	// remove hook from maps
+	s.hooks.Delete(hook)
+	s.hooksOut.Delete(hook)
+	if !hook.expires.IsZero() {
+		s.hookExpires.Delete(hook)
+	}
+	// remove any hook / object connections
+	s.groupDisconnectHook(hook.Name)
+	// remove hook from spatial index
+	if hook.Fence != nil && hook.Fence.obj != nil {
+		rect := hook.Fence.obj.Rect()
+		s.hookTree.Delete(
+			[2]float64{rect.Min.X, rect.Min.Y},
+			[2]float64{rect.Max.X, rect.Max.Y},
+			hook)
+		if hook.Fence.detect["cross"] {
+			s.hookCross.Delete(
+				[2]float64{rect.Min.X, rect.Min.Y},
+				[2]float64{rect.Max.X, rect.Max.Y},
+				hook)
+		}
+	}
+	return true
+}
+
 func (s *Server) cmdDelHook(msg *Message) (
 	res resp.Value, d commandDetails, err error,
 ) {
@@ -255,33 +286,8 @@ func (s *Server) cmdDelHook(msg *Message) (
 	if len(vs) != 0 {
 		return NOMessage, d, errInvalidNumberOfArguments
 	}
-	hook, _ := s.hooks.Get(&Hook{Name: name}).(*Hook)
-	if hook != nil && hook.channel == channel {
-		hook.Close()
-		// remove hook from maps
-		s.hooks.Delete(hook)
-		s.hooksOut.Delete(hook)
-		if !hook.expires.IsZero() {
-			s.hookExpires.Delete(hook)
-		}
-		// remove any hook / object connections
-		s.groupDisconnectHook(hook.Name)
-		// remove hook from spatial index
-		if hook.Fence != nil && hook.Fence.obj != nil {
-			rect := hook.Fence.obj.Rect()
-			s.hookTree.Delete(
-				[2]float64{rect.Min.X, rect.Min.Y},
-				[2]float64{rect.Max.X, rect.Max.Y},
-				hook)
-			if hook.Fence.detect["cross"] {
-				s.hookCross.Delete(
-					[2]float64{rect.Min.X, rect.Min.Y},
-					[2]float64{rect.Max.X, rect.Max.Y},
-					hook)
-			}
-		}
-		d.updated = true
-	}
+
+	d.updated = s.cmdDELHOOKop(name, channel)
 	d.timestamp = time.Now()
 
 	switch msg.OutputType {
@@ -323,29 +329,7 @@ func (s *Server) cmdPDelHook(msg *Message) (
 		if hook.channel != channel {
 			continue
 		}
-		hook.Close()
-		// remove hook from maps
-		s.hooks.Delete(hook)
-		s.hooksOut.Delete(hook)
-		if !hook.expires.IsZero() {
-			s.hookExpires.Delete(hook)
-		}
-		// remove any hook / object connections
-		s.groupDisconnectHook(hook.Name)
-		// remove hook from spatial index
-		if hook.Fence != nil && hook.Fence.obj != nil {
-			rect := hook.Fence.obj.Rect()
-			s.hookTree.Delete(
-				[2]float64{rect.Min.X, rect.Min.Y},
-				[2]float64{rect.Max.X, rect.Max.Y},
-				hook)
-			if hook.Fence.detect["cross"] {
-				s.hookCross.Delete(
-					[2]float64{rect.Min.X, rect.Min.Y},
-					[2]float64{rect.Max.X, rect.Max.Y},
-					hook)
-			}
-		}
+		s.cmdDELHOOKop(hook.Name, channel)
 		d.updated = true
 		count++
 	}