[registry] relocates tag cleanup locking

registry/server/tag_cleanup.go

@@ -27,25 +27,20 @@ func (tc *TagCleaner) RunTagCleanup(s *Server, ctx context.Context, c Config) {
 	for tc.running {
 		<-t.C
 
-		if err := s.Locking.Lock(ctx); err != nil {
-			log.Println(err)
+		if err := s.MarkForDeletion(ctx, time.Now); err != nil {
+			log.Println("error marking tags for deletion: ", err)
 			continue
 		}
-		defer s.Locking.UnlockLogError(ctx)
 
-		err := s.MarkForDeletion(time.Now)
-		if err != nil {
-			log.Println(err)
-			continue
+		if err := s.DeleteStaleTags(ctx, time.Now, c); err != nil {
+			log.Println("error deleting stale tags: ", err)
 		}
-
-		s.DeleteStaleTags(time.Now, c)
 	}
 }
 
 // MarkForDeletion marks stored tags that have been stranded without an associated
 // kafka resource.
-func (s *Server) MarkForDeletion(now func() time.Time) error {
+func (s *Server) MarkForDeletion(ctx context.Context, now func() time.Time) error {
 	markTimeMinutes := fmt.Sprint(now().Unix())
 
 	// Get all brokers from ZK.
@@ -54,6 +49,12 @@ func (s *Server) MarkForDeletion(now func() time.Time) error {
 		return ErrFetchingBrokers
 	}
 
+	// Lock.
+	if err := s.Locking.Lock(ctx); err != nil {
+		return err
+	}
+	defer s.Locking.UnlockLogError(ctx)
+
 	// Get all topics from ZK
 	topics, err := s.ZK.GetTopics([]*regexp.Regexp{topicRegex})
 	topicSet := TopicSetFromSlice(topics)
@@ -110,8 +111,15 @@ func (s *Server) MarkForDeletion(now func() time.Time) error {
 }
 
 // DeleteStaleTags deletes any tags that have not had a kafka resource associated with them.
-func (s *Server) DeleteStaleTags(now func() time.Time, c Config) {
+func (s *Server) DeleteStaleTags(ctx context.Context, now func() time.Time, c Config) error {
 	sweepTime := now().Unix()
+
+	// Lock.
+	if err := s.Locking.Lock(ctx); err != nil {
+		return err
+	}
+	defer s.Locking.UnlockLogError(ctx)
+
 	allTags, _ := s.Tags.Store.GetAllTags()
 
 	for kafkaObject, tags := range allTags {
@@ -132,6 +140,8 @@ func (s *Server) DeleteStaleTags(now func() time.Time, c Config) {
 			log.Printf("deleted tags for non-existent %s %s\n", kafkaObject.Type, kafkaObject.ID)
 		}
 	}
+
+	return nil
 }
 
 // TopicSetFromSlice converts a slice into a TopicSet for convenience

registry/server/tag_cleanup_test.go

@@ -1,10 +1,12 @@
 package server
 
 import (
+	"context"
 	"fmt"
-	"github.com/DataDog/kafka-kit/v3/kafkazk"
 	"testing"
 	"time"
+
+	"github.com/DataDog/kafka-kit/v3/kafkazk"
 )
 
 func TestMarkStaleTags(t *testing.T) {
@@ -21,13 +23,16 @@ func TestMarkStaleTags(t *testing.T) {
 
 	zk := kafkazk.NewZooKeeperStub()
 	th := testTagHandler()
-	s := Server{Tags: th, ZK: zk}
+
+	s := testServer()
+	s.Tags = th
+	s.ZK = zk
 
 	// WHEN
 	th.Store.SetTags(topic, tt)
 	th.Store.SetTags(broker, bt)
 	th.Store.SetTags(noBroker, nbt)
-	s.MarkForDeletion(time.Now)
+	s.MarkForDeletion(context.Background(), time.Now)
 
 	// THEN
 	nbtags, _ := th.Store.GetTags(noBroker)
@@ -54,13 +59,16 @@ func TestDeleteStaleTags(t *testing.T) {
 	bt := TagSet{"foo": "bar", TagMarkTimeKey: fmt.Sprint(markTime.Unix())}
 	broker := KafkaObject{Type: "broker", ID: "not found"}
 
-	th := testTagHandler()
 	zk := kafkazk.NewZooKeeperStub()
-	s := Server{Tags: th, ZK: zk}
+	th := testTagHandler()
+
+	s := testServer()
+	s.Tags = th
+	s.ZK = zk
 
 	//WHEN
 	th.Store.SetTags(broker, bt)
-	s.DeleteStaleTags(func() time.Time { return sweepTime }, Config{TagAllowedStalenessMinutes: 10})
+	s.DeleteStaleTags(context.Background(), func() time.Time { return sweepTime }, Config{TagAllowedStalenessMinutes: 10})
 
 	//THEN
 	btags, _ := th.Store.GetTags(broker)
@@ -77,13 +85,16 @@ func TestUnmarkedTagsAreSafe(t *testing.T) {
 	bt := TagSet{"foo": "bar"}
 	broker := KafkaObject{Type: "broker", ID: "not found"}
 
-	th := testTagHandler()
 	zk := kafkazk.NewZooKeeperStub()
-	s := Server{Tags: th, ZK: zk}
+	th := testTagHandler()
+
+	s := testServer()
+	s.Tags = th
+	s.ZK = zk
 
 	//WHEN
 	th.Store.SetTags(broker, bt)
-	s.DeleteStaleTags(func() time.Time { return sweepTime }, Config{TagAllowedStalenessMinutes: 10})
+	s.DeleteStaleTags(context.Background(), func() time.Time { return sweepTime }, Config{TagAllowedStalenessMinutes: 10})
 
 	//THEN
 	btags, _ := th.Store.GetTags(broker)
@@ -102,12 +113,15 @@ func TestKafkaObjectComesBack(t *testing.T) {
 
 	zk := kafkazk.NewZooKeeperStub()
 	th := testTagHandler()
-	s := Server{Tags: th, ZK: zk}
+
+	s := testServer()
+	s.Tags = th
+	s.ZK = zk
 
 	// WHEN
 	th.Store.SetTags(broker, bt)
 	th.Store.SetTags(topic, tt)
-	s.MarkForDeletion(time.Now)
+	s.MarkForDeletion(context.Background(), time.Now)
 
 	// THEN
 	btags, _ := th.Store.GetTags(broker)