From a872c35b1a9cbd19904010c105281ad15ab687ab Mon Sep 17 00:00:00 2001 From: monkey92t Date: Wed, 28 Dec 2022 22:14:52 +0800 Subject: [PATCH] feat: add ClientName option Signed-off-by: monkey92t --- cluster.go | 11 ++++++++--- cluster_test.go | 19 +++++++++++++++++++ options.go | 6 +++++- options_test.go | 3 +++ redis.go | 4 ++++ redis_test.go | 15 +++++++++++++++ ring.go | 10 +++++++--- ring_test.go | 15 +++++++++++++++ sentinel.go | 11 +++++++++-- sentinel_test.go | 24 ++++++++++++++++++++++++ universal.go | 18 ++++++++++++------ 11 files changed, 121 insertions(+), 15 deletions(-) diff --git a/cluster.go b/cluster.go index a70ff12..d36b5c3 100644 --- a/cluster.go +++ b/cluster.go @@ -29,6 +29,9 @@ type ClusterOptions struct { // A seed list of host:port addresses of cluster nodes. Addrs []string + // ClientName will execute the `CLIENT SETNAME ClientName` command for each conn. + ClientName string + // NewClient creates a cluster node client with provided name and options. NewClient func(opt *Options) *Client @@ -208,6 +211,7 @@ func setupClusterConn(u *url.URL, host string, o *ClusterOptions) (*ClusterOptio func setupClusterQueryParams(u *url.URL, o *ClusterOptions) (*ClusterOptions, error) { q := queryOptions{q: u.Query()} + o.ClientName = q.string("client_name") o.MaxRedirects = q.int("max_redirects") o.ReadOnly = q.bool("read_only") o.RouteByLatency = q.bool("route_by_latency") @@ -250,8 +254,9 @@ func setupClusterQueryParams(u *url.URL, o *ClusterOptions) (*ClusterOptions, er func (opt *ClusterOptions) clientOptions() *Options { return &Options{ - Dialer: opt.Dialer, - OnConnect: opt.OnConnect, + ClientName: opt.ClientName, + Dialer: opt.Dialer, + OnConnect: opt.OnConnect, Username: opt.Username, Password: opt.Password, @@ -871,7 +876,7 @@ func (c *ClusterClient) Close() error { return c.nodes.Close() } -// Do creates a Cmd from the args and processes the cmd. +// Do create a Cmd from the args and processes the cmd. func (c *ClusterClient) Do(ctx context.Context, args ...interface{}) *Cmd { cmd := NewCmd(ctx, args...) _ = c.Process(ctx, cmd) diff --git a/cluster_test.go b/cluster_test.go index f9dc679..2827d3f 100644 --- a/cluster_test.go +++ b/cluster_test.go @@ -589,6 +589,7 @@ var _ = Describe("ClusterClient", func() { Describe("ClusterClient", func() { BeforeEach(func() { opt = redisClusterOptions() + opt.ClientName = "cluster_hi" client = cluster.newClusterClient(ctx, opt) err := client.ForEachMaster(ctx, func(ctx context.Context, master *redis.Client) error { @@ -679,6 +680,20 @@ var _ = Describe("ClusterClient", func() { Expect(assertSlotsEqual(res, wanted)).NotTo(HaveOccurred()) }) + It("should cluster client setname", func() { + err := client.ForEachShard(ctx, func(ctx context.Context, c *redis.Client) error { + return c.Ping(ctx).Err() + }) + Expect(err).NotTo(HaveOccurred()) + + _ = client.ForEachShard(ctx, func(ctx context.Context, c *redis.Client) error { + val, err := c.ClientList(ctx).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).Should(ContainSubstring("name=cluster_hi")) + return nil + }) + }) + It("should CLUSTER NODES", func() { res, err := client.ClusterNodes(ctx).Result() Expect(err).NotTo(HaveOccurred()) @@ -1408,6 +1423,10 @@ func TestParseClusterURL(t *testing.T) { test: "UseDefault", url: "redis://localhost:123?conn_max_idle_time=", o: &redis.ClusterOptions{Addrs: []string{"localhost:123"}, ConnMaxIdleTime: 0}, + }, { + test: "ClientName", + url: "redis://localhost:123?client_name=cluster_hi", + o: &redis.ClusterOptions{Addrs: []string{"localhost:123"}, ClientName: "cluster_hi"}, }, { test: "UseDefaultMissing=", url: "redis://localhost:123?conn_max_idle_time", diff --git a/options.go b/options.go index 651d0d4..49abe21 100644 --- a/options.go +++ b/options.go @@ -27,7 +27,7 @@ type Limiter interface { ReportResult(result error) } -// Options keeps the settings to setup redis connection. +// Options keeps the settings to set up redis connection. type Options struct { // The network type, either tcp or unix. // Default is tcp. @@ -35,6 +35,9 @@ type Options struct { // host:port address. Addr string + // ClientName will execute the `CLIENT SETNAME ClientName` command for each conn. + ClientName string + // Dialer creates new network connection and has priority over // Network and Addr options. Dialer func(ctx context.Context, network, addr string) (net.Conn, error) @@ -426,6 +429,7 @@ func setupConnParams(u *url.URL, o *Options) (*Options, error) { o.DB = db } + o.ClientName = q.string("client_name") o.MaxRetries = q.int("max_retries") o.MinRetryBackoff = q.duration("min_retry_backoff") o.MaxRetryBackoff = q.duration("max_retry_backoff") diff --git a/options_test.go b/options_test.go index 47fa9cb..4ad9175 100644 --- a/options_test.go +++ b/options_test.go @@ -59,6 +59,9 @@ func TestParseURL(t *testing.T) { }, { url: "redis://localhost:123/?db=2&conn_max_idle_time", // missing "=" at the end o: &Options{Addr: "localhost:123", DB: 2, ConnMaxIdleTime: 0}, + }, { + url: "redis://localhost:123/?db=2&client_name=hi", // client name + o: &Options{Addr: "localhost:123", DB: 2, ClientName: "hi"}, }, { url: "unix:///tmp/redis.sock", o: &Options{Addr: "/tmp/redis.sock"}, diff --git a/redis.go b/redis.go index ee6c9bd..d587c1f 100644 --- a/redis.go +++ b/redis.go @@ -256,6 +256,10 @@ func (c *baseClient) initConn(ctx context.Context, cn *pool.Conn) error { pipe.ReadOnly(ctx) } + if c.opt.ClientName != "" { + pipe.ClientSetName(ctx, c.opt.ClientName) + } + return nil }) if err != nil { diff --git a/redis_test.go b/redis_test.go index 8df1dd8..4cbc389 100644 --- a/redis_test.go +++ b/redis_test.go @@ -169,6 +169,21 @@ var _ = Describe("Client", func() { Expect(db2.Close()).NotTo(HaveOccurred()) }) + It("should client setname", func() { + opt := redisOptions() + opt.ClientName = "hi" + db := redis.NewClient(opt) + + defer func() { + Expect(db.Close()).NotTo(HaveOccurred()) + }() + + Expect(db.Ping(ctx).Err()).NotTo(HaveOccurred()) + val, err := db.ClientList(ctx).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).Should(ContainSubstring("name=hi")) + }) + It("processes custom commands", func() { cmd := redis.NewCmd(ctx, "PING") _ = client.Process(ctx, cmd) diff --git a/ring.go b/ring.go index 0a1069d..bc299da 100644 --- a/ring.go +++ b/ring.go @@ -51,6 +51,9 @@ type RingOptions struct { // NewClient creates a shard client with provided options. NewClient func(opt *Options) *Client + // ClientName will execute the `CLIENT SETNAME ClientName` command for each conn. + ClientName string + // Frequency of PING commands sent to check shards availability. // Shard is considered down after 3 subsequent failed checks. HeartbeatFrequency time.Duration @@ -129,8 +132,9 @@ func (opt *RingOptions) init() { func (opt *RingOptions) clientOptions() *Options { return &Options{ - Dialer: opt.Dialer, - OnConnect: opt.OnConnect, + ClientName: opt.ClientName, + Dialer: opt.Dialer, + OnConnect: opt.OnConnect, Username: opt.Username, Password: opt.Password, @@ -522,7 +526,7 @@ func (c *Ring) SetAddrs(addrs map[string]string) { c.sharding.SetAddrs(addrs) } -// Do creates a Cmd from the args and processes the cmd. +// Do create a Cmd from the args and processes the cmd. func (c *Ring) Do(ctx context.Context, args ...interface{}) *Cmd { cmd := NewCmd(ctx, args...) _ = c.Process(ctx, cmd) diff --git a/ring_test.go b/ring_test.go index 7caa868..df33599 100644 --- a/ring_test.go +++ b/ring_test.go @@ -29,6 +29,7 @@ var _ = Describe("Redis Ring", func() { BeforeEach(func() { opt := redisRingOptions() + opt.ClientName = "ring_hi" opt.HeartbeatFrequency = heartbeat ring = redis.NewRing(opt) @@ -50,6 +51,20 @@ var _ = Describe("Redis Ring", func() { Expect(err).To(MatchError("context canceled")) }) + It("should ring client setname", func() { + err := ring.ForEachShard(ctx, func(ctx context.Context, c *redis.Client) error { + return c.Ping(ctx).Err() + }) + Expect(err).NotTo(HaveOccurred()) + + _ = ring.ForEachShard(ctx, func(ctx context.Context, c *redis.Client) error { + val, err := c.ClientList(ctx).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).Should(ContainSubstring("name=ring_hi")) + return nil + }) + }) + It("distributes keys", func() { setRingKeys() diff --git a/sentinel.go b/sentinel.go index 44b073f..1feeb03 100644 --- a/sentinel.go +++ b/sentinel.go @@ -24,6 +24,9 @@ type FailoverOptions struct { // A seed list of host:port addresses of sentinel nodes. SentinelAddrs []string + // ClientName will execute the `CLIENT SETNAME ClientName` command for each conn. + ClientName string + // If specified with SentinelPassword, enables ACL-based authentication (via // AUTH ). SentinelUsername string @@ -78,7 +81,8 @@ type FailoverOptions struct { func (opt *FailoverOptions) clientOptions() *Options { return &Options{ - Addr: "FailoverClient", + Addr: "FailoverClient", + ClientName: opt.ClientName, Dialer: opt.Dialer, OnConnect: opt.OnConnect, @@ -110,7 +114,8 @@ func (opt *FailoverOptions) clientOptions() *Options { func (opt *FailoverOptions) sentinelOptions(addr string) *Options { return &Options{ - Addr: addr, + Addr: addr, + ClientName: opt.ClientName, Dialer: opt.Dialer, OnConnect: opt.OnConnect, @@ -141,6 +146,8 @@ func (opt *FailoverOptions) sentinelOptions(addr string) *Options { func (opt *FailoverOptions) clusterOptions() *ClusterOptions { return &ClusterOptions{ + ClientName: opt.ClientName, + Dialer: opt.Dialer, OnConnect: opt.OnConnect, diff --git a/sentinel_test.go b/sentinel_test.go index ecde161..cfa8b01 100644 --- a/sentinel_test.go +++ b/sentinel_test.go @@ -1,6 +1,7 @@ package redis_test import ( + "context" "net" . "github.com/onsi/ginkgo" @@ -17,6 +18,7 @@ var _ = Describe("Sentinel", func() { BeforeEach(func() { client = redis.NewFailoverClient(&redis.FailoverOptions{ + ClientName: "sentinel_hi", MasterName: sentinelName, SentinelAddrs: sentinelAddrs, MaxRetries: -1, @@ -125,6 +127,13 @@ var _ = Describe("Sentinel", func() { err := client.Ping(ctx).Err() Expect(err).NotTo(HaveOccurred()) }) + + It("should sentinel client setname", func() { + Expect(client.Ping(ctx).Err()).NotTo(HaveOccurred()) + val, err := client.ClientList(ctx).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).Should(ContainSubstring("name=sentinel_hi")) + }) }) var _ = Describe("NewFailoverClusterClient", func() { @@ -134,6 +143,7 @@ var _ = Describe("NewFailoverClusterClient", func() { BeforeEach(func() { client = redis.NewFailoverClusterClient(&redis.FailoverOptions{ + ClientName: "sentinel_cluster_hi", MasterName: sentinelName, SentinelAddrs: sentinelAddrs, @@ -213,6 +223,20 @@ var _ = Describe("NewFailoverClusterClient", func() { _, err = startRedis(masterPort) Expect(err).NotTo(HaveOccurred()) }) + + It("should sentinel cluster client setname", func() { + err := client.ForEachShard(ctx, func(ctx context.Context, c *redis.Client) error { + return c.Ping(ctx).Err() + }) + Expect(err).NotTo(HaveOccurred()) + + _ = client.ForEachShard(ctx, func(ctx context.Context, c *redis.Client) error { + val, err := c.ClientList(ctx).Result() + Expect(err).NotTo(HaveOccurred()) + Expect(val).Should(ContainSubstring("name=sentinel_cluster_hi")) + return nil + }) + }) }) var _ = Describe("SentinelAclAuth", func() { diff --git a/universal.go b/universal.go index 73bbf17..9d1a852 100644 --- a/universal.go +++ b/universal.go @@ -14,6 +14,9 @@ type UniversalOptions struct { // of cluster/sentinel nodes. Addrs []string + // ClientName will execute the `CLIENT SETNAME ClientName` command for each conn. + ClientName string + // Database to be selected after connecting to the server. // Only single-node and failover clients. DB int @@ -69,9 +72,10 @@ func (o *UniversalOptions) Cluster() *ClusterOptions { } return &ClusterOptions{ - Addrs: o.Addrs, - Dialer: o.Dialer, - OnConnect: o.OnConnect, + Addrs: o.Addrs, + ClientName: o.ClientName, + Dialer: o.Dialer, + OnConnect: o.OnConnect, Username: o.Username, Password: o.Password, @@ -112,6 +116,7 @@ func (o *UniversalOptions) Failover() *FailoverOptions { return &FailoverOptions{ SentinelAddrs: o.Addrs, MasterName: o.MasterName, + ClientName: o.ClientName, Dialer: o.Dialer, OnConnect: o.OnConnect, @@ -151,9 +156,10 @@ func (o *UniversalOptions) Simple() *Options { } return &Options{ - Addr: addr, - Dialer: o.Dialer, - OnConnect: o.OnConnect, + Addr: addr, + ClientName: o.ClientName, + Dialer: o.Dialer, + OnConnect: o.OnConnect, DB: o.DB, Username: o.Username,