Adding new commands MExpire/MExpireAt, fixed the Ping command, minor performance improvements

Author: nvorobev

CHANGELOG.md

@@ -1,3 +1,9 @@
+## v2.2.0 (2023-04-05)
+
+- Adding new commands MExpire/MExpireAt with support for executing multiple keys
+- Fixed the Ping command, now the check occurs on all shards and returns the first error for the sharded configuration
+- Minor performance improvements
+
 ## v2.1.0 (2022-07-12)
 
 - Added getting statistics of all connection pools

pool.go

@@ -206,12 +206,20 @@ func (p *Pool) TxPipelined(ctx context.Context, fn func(redis.Pipeliner) error)
 }
 
 func (p *Pool) Ping(ctx context.Context) *redis.StatusCmd {
-	// FIXME: use config to determine whether no key would access the master
-	conn, err := p.connFactory.getMasterConn()
-	if err != nil {
-		return newErrorStatusCmd(err)
+	if _, ok := p.connFactory.(*HAConnFactory); ok {
+		conn, _ := p.connFactory.getMasterConn()
+		return conn.Ping(ctx)
+	}
+	var result *redis.StatusCmd
+	factory := p.connFactory.(*ShardConnFactory)
+	for _, shard := range factory.shards {
+		conn, _ := shard.getMasterConn()
+		result = conn.Ping(ctx)
+		if result.Err() != nil {
+			return result
+		}
 	}
-	return conn.Ping(ctx)
+	return result
 }
 
 func (p *Pool) Get(ctx context.Context, key string) *redis.StringCmd {
@@ -425,18 +433,17 @@ func (p *Pool) MSetWithGD(ctx context.Context, values ...interface{}) []*redis.S
 	}
 
 	var wg sync.WaitGroup
-	var mu sync.Mutex
-	var result []*redis.StatusCmd
+	result := make([]*redis.StatusCmd, len(index2Values))
+	var i int
 	for ind, vals := range index2Values {
 		wg.Add(1)
 		conn, _ := factory.shards[ind].getMasterConn()
-		go func(conn *redis.Client, vals ...interface{}) {
-			defer wg.Done()
+		go func(i int, conn *redis.Client, vals ...interface{}) {
 			status := conn.MSet(ctx, vals...)
-			mu.Lock()
-			result = append(result, status)
-			mu.Unlock()
-		}(conn, vals...)
+			result[i] = status
+			wg.Done()
+		}(i, conn, vals...)
+		i++
 	}
 	wg.Wait()
 	return result
@@ -519,6 +526,79 @@ func (p *Pool) Expire(ctx context.Context, key string, expiration time.Duration)
 	return conn.Expire(ctx, key, expiration)
 }
 
+// MExpire gives the result for each group of keys
+func (p *Pool) MExpire(ctx context.Context, expiration time.Duration, keys ...string) map[string]error {
+	keyErrorsMap := func(results []redis.Cmder) map[string]error {
+		if len(results) == 0 {
+			return nil
+		}
+		keyErrors := make(map[string]error, 0)
+		for _, result := range results {
+			if result.Err() != nil {
+				args := result.Args()
+				for i, arg := range args {
+					if i == 0 || i == 2 {
+						continue
+					}
+					keyErrors[arg.(string)] = result.Err()
+				}
+			}
+		}
+		return keyErrors
+	}
+
+	if _, ok := p.connFactory.(*HAConnFactory); ok {
+		conn, _ := p.connFactory.getMasterConn()
+		pipe := conn.Pipeline()
+		for _, key := range keys {
+			pipe.Expire(ctx, key, expiration)
+		}
+		results, err := pipe.Exec(ctx)
+		_ = pipe.Close()
+		if err != nil {
+			return keyErrorsMap(results)
+		}
+		return nil
+	}
+
+	factory := p.connFactory.(*ShardConnFactory)
+	index2Keys := make(map[uint32][]string)
+	for _, key := range keys {
+		ind := factory.getShardIndex(key)
+		if _, ok := index2Keys[ind]; !ok {
+			index2Keys[ind] = make([]string, 0)
+		}
+		index2Keys[ind] = append(index2Keys[ind], key)
+	}
+
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+	var results []redis.Cmder
+	var i int
+	for ind, keys := range index2Keys {
+		wg.Add(1)
+		conn, _ := factory.shards[ind].getMasterConn()
+		go func(i int, conn *redis.Client, keys ...string) {
+			pipe := conn.Pipeline()
+			for _, key := range keys {
+				pipe.Expire(ctx, key, expiration)
+			}
+			result, err := pipe.Exec(ctx)
+			_ = pipe.Close()
+			if err != nil {
+				mu.Lock()
+				results = append(results, result...)
+				mu.Unlock()
+			}
+			wg.Done()
+		}(i, conn, keys...)
+		i++
+	}
+	wg.Wait()
+
+	return keyErrorsMap(results)
+}
+
 func (p *Pool) ExpireAt(ctx context.Context, key string, tm time.Time) *redis.BoolCmd {
 	conn, err := p.connFactory.getMasterConn(key)
 	if err != nil {
@@ -527,6 +607,79 @@ func (p *Pool) ExpireAt(ctx context.Context, key string, tm time.Time) *redis.Bo
 	return conn.ExpireAt(ctx, key, tm)
 }
 
+// MExpireAt gives the result for each group of keys
+func (p *Pool) MExpireAt(ctx context.Context, tm time.Time, keys ...string) map[string]error {
+	keyErrorsMap := func(results []redis.Cmder) map[string]error {
+		if len(results) == 0 {
+			return nil
+		}
+		keyErrors := make(map[string]error, 0)
+		for _, result := range results {
+			if result.Err() != nil {
+				args := result.Args()
+				for i, arg := range args {
+					if i == 0 || i == 2 {
+						continue
+					}
+					keyErrors[arg.(string)] = result.Err()
+				}
+			}
+		}
+		return keyErrors
+	}
+
+	if _, ok := p.connFactory.(*HAConnFactory); ok {
+		conn, _ := p.connFactory.getMasterConn()
+		pipe := conn.Pipeline()
+		for _, key := range keys {
+			pipe.ExpireAt(ctx, key, tm)
+		}
+		results, err := pipe.Exec(ctx)
+		_ = pipe.Close()
+		if err != nil {
+			return keyErrorsMap(results)
+		}
+		return nil
+	}
+
+	factory := p.connFactory.(*ShardConnFactory)
+	index2Keys := make(map[uint32][]string)
+	for _, key := range keys {
+		ind := factory.getShardIndex(key)
+		if _, ok := index2Keys[ind]; !ok {
+			index2Keys[ind] = make([]string, 0)
+		}
+		index2Keys[ind] = append(index2Keys[ind], key)
+	}
+
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+	var results []redis.Cmder
+	var i int
+	for ind, keys := range index2Keys {
+		wg.Add(1)
+		conn, _ := factory.shards[ind].getMasterConn()
+		go func(i int, conn *redis.Client, keys ...string) {
+			pipe := conn.Pipeline()
+			for _, key := range keys {
+				pipe.ExpireAt(ctx, key, tm)
+			}
+			result, err := pipe.Exec(ctx)
+			_ = pipe.Close()
+			if err != nil {
+				mu.Lock()
+				results = append(results, result...)
+				mu.Unlock()
+			}
+			wg.Done()
+		}(i, conn, keys...)
+		i++
+	}
+	wg.Wait()
+
+	return keyErrorsMap(results)
+}
+
 func (p *Pool) TTL(ctx context.Context, key string) *redis.DurationCmd {
 	conn, err := p.connFactory.getSlaveConn(key)
 	if err != nil {

pool_test.go

@@ -1820,6 +1820,14 @@ var _ = Describe("Pool_GD", func() {
 	})
 
 	Describe("Commands", func() {
+		It("ping", func() {
+			for _, pool := range pools {
+				_, err := pool.Ping(ctx).Result()
+				Expect(err).To(HaveOccurred())
+				Expect(err.Error()).To(Equal("dial tcp 127.0.0.1:8384: connect: connection refused"))
+			}
+		})
+
 		It("gd", func() {
 			kvs := []string{"a3", "a3", "b3", "b3", "c3", "c3", "d3", "d3"}
 			keys := make([]string, 0)
@@ -1844,5 +1852,53 @@ var _ = Describe("Pool_GD", func() {
 				_, _ = pool.Del(ctx, keys...)
 			}
 		})
+
+		It("mexpire", func() {
+			kvs := []string{"a3", "a3", "b3", "b3", "c3", "c3", "d3", "d3"}
+			keys := make([]string, 0)
+			for i := 0; i < len(kvs); i += 2 {
+				keys = append(keys, kvs[i])
+			}
+			for _, pool := range pools {
+				statuses := pool.MSetWithGD(ctx, kvs)
+				time.Sleep(10 * time.Millisecond)
+				sort.Slice(statuses, func(i, j int) bool {
+					return statuses[i].Err() == nil
+				})
+				Expect(statuses[0].Err()).NotTo(HaveOccurred())
+				Expect(statuses[1].Err()).To(HaveOccurred())
+				keyErrors := pool.MExpire(ctx, 5*time.Minute, keys...)
+				time.Sleep(10 * time.Millisecond)
+				Expect(keyErrors).Should(HaveKey("a3"))
+				Expect(keyErrors).Should(HaveKey("c3"))
+				Expect(pool.TTL(ctx, "b3").Val()).NotTo(Equal(-1))
+				Expect(pool.TTL(ctx, "d3").Val()).NotTo(Equal(-1))
+				_, _ = pool.Del(ctx, keys...)
+			}
+		})
+
+		It("mexpireat", func() {
+			kvs := []string{"a3", "a3", "b3", "b3", "c3", "c3", "d3", "d3"}
+			keys := make([]string, 0)
+			for i := 0; i < len(kvs); i += 2 {
+				keys = append(keys, kvs[i])
+			}
+			for _, pool := range pools {
+				statuses := pool.MSetWithGD(ctx, kvs)
+				time.Sleep(10 * time.Millisecond)
+				sort.Slice(statuses, func(i, j int) bool {
+					return statuses[i].Err() == nil
+				})
+				Expect(statuses[0].Err()).NotTo(HaveOccurred())
+				Expect(statuses[1].Err()).To(HaveOccurred())
+				keyErrors := pool.MExpireAt(ctx, time.Now().Add(5*time.Minute), keys...)
+				time.Sleep(10 * time.Millisecond)
+				Expect(keyErrors).Should(HaveKey("a3"))
+				Expect(keyErrors).Should(HaveKey("c3"))
+				Expect(pool.TTL(ctx, "b3").Val()).NotTo(Equal(-1))
+				Expect(pool.TTL(ctx, "d3").Val()).NotTo(Equal(-1))
+				_, _ = pool.Del(ctx, keys...)
+			}
+		})
 	})
 })

shard_conn_factory.go

@@ -65,14 +65,12 @@ func NewShardConnFactory(cfg *ShardConfig) (*ShardConnFactory, error) {
 
 func (factory *ShardConnFactory) stats() map[string]*redis.PoolStats {
 	results := make(map[string]*redis.PoolStats, len(factory.shards))
-
 	for _, shard := range factory.shards {
 		result := shard.stats()
 		for addr, stats := range result {
 			results[addr] = stats
 		}
 	}
-
 	return results
 }
 
@@ -152,18 +150,17 @@ func (factory *ShardConnFactory) doMultiKeys(fn multiKeyFn, keys ...string) []re
 		return []redis.Cmder{fn(factory, keys...)}
 	}
 
-	var mu sync.Mutex
 	var wg sync.WaitGroup
-	var results []redis.Cmder
+	results := make([]redis.Cmder, len(index2Keys))
+	var i int
 	for _, keyList := range index2Keys {
 		wg.Add(1)
-		go func(keyList []string) {
-			defer wg.Done()
+		go func(keyList []string, i int) {
 			result := fn(factory, keyList...)
-			mu.Lock()
-			results = append(results, result)
-			mu.Unlock()
-		}(keyList)
+			results[i] = result
+			wg.Done()
+		}(keyList, i)
+		i++
 	}
 	wg.Wait()
 	return results