pubsub pool

Author: ndyakov

example/pubsub/main.go

@@ -11,12 +11,15 @@ import (
 )
 
 var ctx = context.Background()
-var consStopped = false
 
+// This example is not supposed to be run as is. It is just a test to see how pubsub behaves in relation to pool management.
+// It was used to find regressions in pool management in hitless mode.
+// Please don't use it as a reference for how to use pubsub.
 func main() {
 	wg := &sync.WaitGroup{}
 	rdb := redis.NewClient(&redis.Options{
-		Addr: ":6379",
+		Addr:            ":6379",
+		HitlessUpgrades: true,
 	})
 	_ = rdb.FlushDB(ctx).Err()
 
@@ -30,21 +33,22 @@ func main() {
 	if err != nil {
 		panic(err)
 	}
-	if err := rdb.Set(ctx, "prods", "0", 0).Err(); err != nil {
+	if err := rdb.Set(ctx, "publishers", "0", 0).Err(); err != nil {
 		panic(err)
 	}
-	if err := rdb.Set(ctx, "cons", "0", 0).Err(); err != nil {
+	if err := rdb.Set(ctx, "subscribers", "0", 0).Err(); err != nil {
 		panic(err)
 	}
-	if err := rdb.Set(ctx, "cntr", "0", 0).Err(); err != nil {
+	if err := rdb.Set(ctx, "published", "0", 0).Err(); err != nil {
 		panic(err)
 	}
-	if err := rdb.Set(ctx, "recs", "0", 0).Err(); err != nil {
+	if err := rdb.Set(ctx, "received", "0", 0).Err(); err != nil {
 		panic(err)
 	}
-	fmt.Println("cntr", rdb.Get(ctx, "cntr").Val())
-	fmt.Println("recs", rdb.Get(ctx, "recs").Val())
+	fmt.Println("published", rdb.Get(ctx, "published").Val())
+	fmt.Println("received", rdb.Get(ctx, "received").Val())
 	subCtx, cancelSubCtx := context.WithCancel(ctx)
+	pubCtx, cancelPublishers := context.WithCancel(ctx)
 	for i := 0; i < 10; i++ {
 		wg.Add(1)
 		go subscribe(subCtx, rdb, "test", i, wg)
@@ -54,32 +58,39 @@ func main() {
 	time.Sleep(time.Second)
 	subCtx, cancelSubCtx = context.WithCancel(ctx)
 	for i := 0; i < 10; i++ {
-		if err := rdb.Incr(ctx, "prods").Err(); err != nil {
+		if err := rdb.Incr(ctx, "publishers").Err(); err != nil {
 			panic(err)
 		}
 		wg.Add(1)
-		go floodThePool(subCtx, rdb, wg)
+		go floodThePool(pubCtx, rdb, wg)
 	}
 
 	for i := 0; i < 500; i++ {
-		if err := rdb.Incr(ctx, "cons").Err(); err != nil {
+		if err := rdb.Incr(ctx, "subscribers").Err(); err != nil {
 			panic(err)
 		}
 		wg.Add(1)
 		go subscribe(subCtx, rdb, "test2", i, wg)
 	}
+	time.Sleep(5 * time.Second)
+	fmt.Println("canceling publishers")
+	cancelPublishers()
 	time.Sleep(10 * time.Second)
-	fmt.Println("canceling")
+	fmt.Println("canceling subscribers")
 	cancelSubCtx()
 	wg.Wait()
-	cntr, err := rdb.Get(ctx, "cntr").Result()
-	recs, err := rdb.Get(ctx, "recs").Result()
-	prods, err := rdb.Get(ctx, "prods").Result()
-	cons, err := rdb.Get(ctx, "cons").Result()
-	fmt.Printf("cntr: %s\n", cntr)
-	fmt.Printf("recs: %s\n", recs)
-	fmt.Printf("prods: %s\n", prods)
-	fmt.Printf("cons: %s\n", cons)
+	published, err := rdb.Get(ctx, "published").Result()
+	received, err := rdb.Get(ctx, "received").Result()
+	publishers, err := rdb.Get(ctx, "publishers").Result()
+	subscribers, err := rdb.Get(ctx, "subscribers").Result()
+	fmt.Printf("publishers: %s\n", publishers)
+	fmt.Printf("published: %s\n", published)
+	fmt.Printf("subscribers: %s\n", subscribers)
+	fmt.Printf("received: %s\n", received)
+	publishedInt, err := rdb.Get(ctx, "published").Int()
+	subscribersInt, err := rdb.Get(ctx, "subscribers").Int()
+	fmt.Printf("if drained = published*subscribers: %d\n", publishedInt*subscribersInt)
+
 	time.Sleep(2 * time.Second)
 }
 
@@ -88,8 +99,6 @@ func floodThePool(ctx context.Context, rdb *redis.Client, wg *sync.WaitGroup) {
 	for {
 		select {
 		case <-ctx.Done():
-			fmt.Println("floodThePool stopping")
-			consStopped = true
 			return
 		default:
 		}
@@ -99,7 +108,7 @@ func floodThePool(ctx context.Context, rdb *redis.Client, wg *sync.WaitGroup) {
 			//log.Println("publish error:", err)
 		}
 
-		err = rdb.Incr(ctx, "cntr").Err()
+		err = rdb.Incr(ctx, "published").Err()
 		if err != nil {
 			// noop
 			//log.Println("incr error:", err)
@@ -110,36 +119,24 @@ func floodThePool(ctx context.Context, rdb *redis.Client, wg *sync.WaitGroup) {
 
 func subscribe(ctx context.Context, rdb *redis.Client, topic string, subscriberId int, wg *sync.WaitGroup) {
 	defer wg.Done()
-	defer fmt.Printf("subscriber %d stopping\n", subscriberId)
 	rec := rdb.Subscribe(ctx, topic)
 	recChan := rec.Channel()
 	for {
 		select {
 		case <-ctx.Done():
 			rec.Close()
-			if subscriberId == 199 {
-				fmt.Printf("subscriber %d done\n", subscriberId)
-			}
 			return
 		default:
 			select {
 			case <-ctx.Done():
 				rec.Close()
-				if subscriberId == 199 {
-					fmt.Printf("subscriber %d done\n", subscriberId)
-				}
 				return
 			case msg := <-recChan:
-				err := rdb.Incr(ctx, "recs").Err()
+				err := rdb.Incr(ctx, "received").Err()
 				if err != nil {
 					log.Println("incr error:", err)
 				}
-				if consStopped {
-					fmt.Printf("subscriber %d received %s\n", subscriberId, msg.Payload)
-				}
-				if subscriberId == 199 {
-					fmt.Printf("subscriber %d received %s\n", subscriberId, msg.Payload)
-				}
+				_ = msg // Use the message to avoid unused variable warning
 			}
 		}
 	}

hitless/config.go

@@ -314,13 +314,14 @@ func isPrivateIP(addr string) bool {
 	// Simplified check for common private IP ranges
 	// 10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16
 	// This is a simplified implementation; a full implementation would parse the IP properly
-	if len(addr) >= 3 {
-		if addr[:3] == "10." || addr[:8] == "192.168." {
-			return true
-		}
-		if len(addr) >= 7 && addr[:7] == "172.16." {
-			return true
-		}
+	if len(addr) >= 3 && addr[:3] == "10." {
+		return true
+	}
+	if len(addr) >= 8 && addr[:8] == "192.168." {
+		return true
+	}
+	if len(addr) >= 7 && addr[:7] == "172.16." {
+		return true
 	}
 	return false
 }

hitless/redis_connection_processor_test.go

@@ -60,10 +60,6 @@ func (mp *mockPool) Get(ctx context.Context) (*pool.Conn, error) {
 	return nil, errors.New("not implemented")
 }
 
-func (mp *mockPool) GetPubSub(ctx context.Context) (*pool.Conn, error) {
-	return nil, errors.New("not implemented")
-}
-
 func (mp *mockPool) Put(ctx context.Context, conn *pool.Conn) {
 	// Not implemented for testing
 }
@@ -115,10 +111,13 @@ func TestRedisConnectionProcessor(t *testing.T) {
 			t.Fatalf("Failed to mark connection for handoff: %v", err)
 		}
 
-		// Set a mock initialization function
-		initConnCalled := false
+		// Set a mock initialization function with synchronization
+		initConnCalled := make(chan bool, 1)
 		initConnFunc := func(ctx context.Context, cn *pool.Conn) error {
-			initConnCalled = true
+			select {
+			case initConnCalled <- true:
+			default:
+			}
 			return nil
 		}
 		conn.SetInitConnFunc(initConnFunc)
@@ -142,22 +141,44 @@ func TestRedisConnectionProcessor(t *testing.T) {
 			t.Error("Connection should be in pending handoffs map")
 		}
 
-		// Wait for handoff to complete
-		time.Sleep(100 * time.Millisecond)
+		// Wait for initialization to be called (indicates handoff started)
+		select {
+		case <-initConnCalled:
+			// Good, initialization was called
+		case <-time.After(1 * time.Second):
+			t.Fatal("Timeout waiting for initialization function to be called")
+		}
+
+		// Wait for handoff to complete with proper timeout and polling
+		timeout := time.After(2 * time.Second)
+		ticker := time.NewTicker(10 * time.Millisecond)
+		defer ticker.Stop()
+
+		handoffCompleted := false
+		for !handoffCompleted {
+			select {
+			case <-timeout:
+				t.Fatal("Timeout waiting for handoff to complete")
+			case <-ticker.C:
+				if _, pending := processor.pending.Load(conn); !pending {
+					handoffCompleted = true
+				}
+			}
+		}
 
 		// Verify handoff completed (removed from pending map)
 		if _, pending := processor.pending.Load(conn); pending {
 			t.Error("Connection should be removed from pending map after handoff")
 		}
 
-		// Verify handoff state was cleared
-		if conn.ShouldHandoff() {
-			t.Error("Connection should not be marked for handoff after successful handoff")
+		// Verify connection is usable again
+		if !conn.IsUsable() {
+			t.Error("Connection should be usable after successful handoff")
 		}
 
-		// Verify initialization was called
-		if !initConnCalled {
-			t.Error("InitConn should have been called")
+		// Verify handoff state is cleared
+		if conn.ShouldHandoff() {
+			t.Error("Connection should not be marked for handoff after completion")
 		}
 	})
 
@@ -236,8 +257,22 @@ func TestRedisConnectionProcessor(t *testing.T) {
 			t.Error("Connection should not be removed when queuing handoff")
 		}
 
-		// Wait for handoff to complete and fail
-		time.Sleep(100 * time.Millisecond)
+		// Wait for handoff to complete and fail with proper timeout and polling
+		timeout := time.After(2 * time.Second)
+		ticker := time.NewTicker(10 * time.Millisecond)
+		defer ticker.Stop()
+
+		handoffCompleted := false
+		for !handoffCompleted {
+			select {
+			case <-timeout:
+				t.Fatal("Timeout waiting for failed handoff to complete")
+			case <-ticker.C:
+				if _, pending := processor.pending.Load(conn); !pending {
+					handoffCompleted = true
+				}
+			}
+		}
 
 		// Connection should be removed from pending map after failed handoff
 		if _, pending := processor.pending.Load(conn); pending {
@@ -468,8 +503,22 @@ func TestRedisConnectionProcessor(t *testing.T) {
 			t.Error("Connection should be pooled after handoff")
 		}
 
-		// Wait for handoff to complete
-		time.Sleep(50 * time.Millisecond)
+		// Wait for handoff to complete with proper timeout and polling
+		timeout := time.After(1 * time.Second)
+		ticker := time.NewTicker(5 * time.Millisecond)
+		defer ticker.Stop()
+
+		handoffCompleted := false
+		for !handoffCompleted {
+			select {
+			case <-timeout:
+				t.Fatal("Timeout waiting for handoff to complete")
+			case <-ticker.C:
+				if _, pending := processor.pending.Load(conn); !pending {
+					handoffCompleted = true
+				}
+			}
+		}
 
 		// Verify relaxed timeout is set with deadline
 		if !conn.HasRelaxedTimeout() {
@@ -626,17 +675,15 @@ func TestRedisConnectionProcessor(t *testing.T) {
 			}
 		}
 
-		// Verify queue has items but capacity remains static
-		currentQueueSize := len(processor.handoffQueue)
-		if currentQueueSize == 0 {
-			t.Error("Expected some items in queue after processing connections")
-		}
-
+		// Verify queue capacity remains static (the main purpose of this test)
 		finalCapacity := cap(processor.handoffQueue)
 		if finalCapacity != 50 {
 			t.Errorf("Queue capacity should remain static at 50, got %d", finalCapacity)
 		}
 
+		// Note: We don't check queue size here because workers process items quickly
+		// The important thing is that the capacity remains static regardless of pool size
+		currentQueueSize := len(processor.handoffQueue)
 		t.Logf("Static queue test completed - Capacity: %d, Current size: %d",
 			finalCapacity, currentQueueSize)
 	})
@@ -738,7 +785,21 @@ func TestRedisConnectionProcessor(t *testing.T) {
 		}
 
 		// Wait for the handoff to complete (it happens asynchronously)
-		time.Sleep(50 * time.Millisecond)
+		timeout := time.After(1 * time.Second)
+		ticker := time.NewTicker(5 * time.Millisecond)
+		defer ticker.Stop()
+
+		handoffCompleted := false
+		for !handoffCompleted {
+			select {
+			case <-timeout:
+				t.Fatal("Timeout waiting for handoff to complete")
+			case <-ticker.C:
+				if _, pending := processor.pending.Load(conn); !pending {
+					handoffCompleted = true
+				}
+			}
+		}
 
 		// Verify that relaxed timeout was applied to the new connection
 		if !conn.HasRelaxedTimeout() {

internal/pool/conn.go

@@ -45,7 +45,6 @@ type Conn struct {
 
 	Inited    bool
 	pooled    bool
-	isPubSub  bool
 	createdAt time.Time
 	expiresAt time.Time