Skip to content

Commit 12b77d7

Browse files
smallnestsmallnest
authored
Merge pull request #914 from baisiyi/feature/xclientpool-plugin-support
feat: support SetPlugins and GetPlugins for XClientPool and OneClientPool
2 parents d563e91 + e389a5f commit 12b77d7

File tree

4 files changed

+300
-0
lines changed

4 files changed

+300
-0
lines changed

โ€Žclient/oneclient_pool.goโ€Ž

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ type OneClientPool struct {
1414
index uint64
1515
oneclients []*OneClient
1616
auth string
17+
Plugins PluginContainer
1718

1819
failMode FailMode
1920
selectMode SelectMode
@@ -68,6 +69,20 @@ func (p *OneClientPool) Auth(auth string) {
6869
}
6970
}
7071

72+
// SetPlugins sets client's plugins.
73+
func (p *OneClientPool) SetPlugins(plugins PluginContainer) {
74+
p.Plugins = plugins
75+
76+
for _, v := range p.oneclients {
77+
v.SetPlugins(plugins)
78+
}
79+
}
80+
81+
// GetPlugins returns client's plugins.
82+
func (p *OneClientPool) GetPlugins() PluginContainer {
83+
return p.Plugins
84+
}
85+
7186
// Get returns a OneClient.
7287
// It does not remove this OneClient from its cache so you don't need to put it back.
7388
// Don't close this OneClient because maybe other goroutines are using this OneClient.

โ€Žclient/oneclient_pool_test.goโ€Ž

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
package client
2+
3+
import (
4+
"testing"
5+
)
6+
7+
func TestOneClientPool_SetPlugins(t *testing.T) {
8+
// Create a simple discovery
9+
pairs := []*KVPair{
10+
{Key: "[email protected]:8972", Value: ""},
11+
}
12+
discovery, err := NewMultipleServersDiscovery(pairs)
13+
if err != nil {
14+
t.Fatalf("failed to create discovery: %v", err)
15+
}
16+
defer discovery.Close()
17+
18+
// Create a pool
19+
pool := NewOneClientPool(3, Failtry, RandomSelect, discovery, DefaultOption)
20+
defer pool.Close()
21+
22+
// Create plugins
23+
plugins := NewPluginContainer()
24+
tp := &testPlugin{name: "test-plugin"}
25+
plugins.Add(tp)
26+
27+
// Test SetPlugins
28+
pool.SetPlugins(plugins)
29+
30+
// Verify plugins are set on pool
31+
if pool.GetPlugins() == nil {
32+
t.Error("plugins should not be nil after SetPlugins")
33+
}
34+
if pool.GetPlugins() != plugins {
35+
t.Error("pool plugins should be the same as the set plugins")
36+
}
37+
38+
// Verify plugins are set on all oneclients
39+
for i := 0; i < 3; i++ {
40+
oneclient := pool.Get()
41+
if oneclient.GetPlugins() == nil {
42+
t.Errorf("oneclient %d plugins should not be nil", i)
43+
}
44+
if oneclient.GetPlugins() != plugins {
45+
t.Errorf("oneclient %d plugins should be the same as the set plugins", i)
46+
}
47+
}
48+
}
49+
50+
func TestOneClientPool_GetPlugins(t *testing.T) {
51+
// Create a simple discovery
52+
pairs := []*KVPair{
53+
{Key: "[email protected]:8972", Value: ""},
54+
}
55+
discovery, err := NewMultipleServersDiscovery(pairs)
56+
if err != nil {
57+
t.Fatalf("failed to create discovery: %v", err)
58+
}
59+
defer discovery.Close()
60+
61+
// Create a pool
62+
pool := NewOneClientPool(2, Failtry, RandomSelect, discovery, DefaultOption)
63+
defer pool.Close()
64+
65+
// Initially, plugins should be nil
66+
if pool.GetPlugins() != nil {
67+
t.Error("plugins should be nil initially")
68+
}
69+
70+
// Create and set plugins
71+
plugins := NewPluginContainer()
72+
tp := &testPlugin{name: "test-plugin"}
73+
plugins.Add(tp)
74+
pool.SetPlugins(plugins)
75+
76+
// Verify GetPlugins returns the correct plugins
77+
retrievedPlugins := pool.GetPlugins()
78+
if retrievedPlugins == nil {
79+
t.Error("plugins should not be nil after SetPlugins")
80+
}
81+
if retrievedPlugins != plugins {
82+
t.Error("GetPlugins should return the same plugins as SetPlugins")
83+
}
84+
85+
// Verify plugins contain the test plugin
86+
allPlugins := retrievedPlugins.All()
87+
if len(allPlugins) != 1 {
88+
t.Errorf("expected 1 plugin, got %d", len(allPlugins))
89+
}
90+
if p, ok := allPlugins[0].(*testPlugin); !ok || p.name != "test-plugin" {
91+
t.Error("plugin should be the test plugin")
92+
}
93+
}
94+
95+
func TestOneClientPool_SetPlugins_Concurrent(t *testing.T) {
96+
// Create a simple discovery
97+
pairs := []*KVPair{
98+
{Key: "[email protected]:8972", Value: ""},
99+
}
100+
discovery, err := NewMultipleServersDiscovery(pairs)
101+
if err != nil {
102+
t.Fatalf("failed to create discovery: %v", err)
103+
}
104+
defer discovery.Close()
105+
106+
// Create a pool
107+
pool := NewOneClientPool(5, Failtry, RandomSelect, discovery, DefaultOption)
108+
defer pool.Close()
109+
110+
// Test concurrent SetPlugins calls
111+
done := make(chan bool, 10)
112+
for i := 0; i < 10; i++ {
113+
go func(id int) {
114+
plugins := NewPluginContainer()
115+
tp := &testPlugin{name: "test-plugin"}
116+
plugins.Add(tp)
117+
pool.SetPlugins(plugins)
118+
done <- true
119+
}(i)
120+
}
121+
122+
// Wait for all goroutines to complete
123+
for i := 0; i < 10; i++ {
124+
<-done
125+
}
126+
127+
// Verify plugins are set
128+
if pool.GetPlugins() == nil {
129+
t.Error("plugins should not be nil after concurrent SetPlugins")
130+
}
131+
}
132+

โ€Žclient/xclient_pool.goโ€Ž

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ type XClientPool struct {
2222
discovery ServiceDiscovery
2323
option Option
2424
auth string
25+
Plugins PluginContainer
2526

2627
serverMessageChan chan<- *protocol.Message
2728
}
@@ -75,6 +76,21 @@ func (c *XClientPool) Auth(auth string) {
7576
c.mu.RUnlock()
7677
}
7778

79+
// SetPlugins sets client's plugins.
80+
func (p *XClientPool) SetPlugins(plugins PluginContainer) {
81+
p.Plugins = plugins
82+
p.mu.RLock()
83+
for _, v := range p.xclients {
84+
v.SetPlugins(plugins)
85+
}
86+
p.mu.RUnlock()
87+
}
88+
89+
// GetPlugins returns client's plugins.
90+
func (p *XClientPool) GetPlugins() PluginContainer {
91+
return p.Plugins
92+
}
93+
7894
// Get returns a xclient.
7995
// It does not remove this xclient from its cache so you don't need to put it back.
8096
// Don't close this xclient because maybe other goroutines are using this xclient.

โ€Žclient/xclient_pool_test.goโ€Ž

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
package client
2+
3+
import (
4+
"testing"
5+
)
6+
7+
// testPlugin is a simple test plugin implementation
8+
type testPlugin struct {
9+
name string
10+
}
11+
12+
func TestXClientPool_SetPlugins(t *testing.T) {
13+
// Create a simple discovery
14+
pairs := []*KVPair{
15+
{Key: "[email protected]:8972", Value: ""},
16+
}
17+
discovery, err := NewMultipleServersDiscovery(pairs)
18+
if err != nil {
19+
t.Fatalf("failed to create discovery: %v", err)
20+
}
21+
defer discovery.Close()
22+
23+
// Create a pool
24+
pool := NewXClientPool(3, "Arith", Failtry, RandomSelect, discovery, DefaultOption)
25+
defer pool.Close()
26+
27+
// Create plugins
28+
plugins := NewPluginContainer()
29+
tp := &testPlugin{name: "test-plugin"}
30+
plugins.Add(tp)
31+
32+
// Test SetPlugins
33+
pool.SetPlugins(plugins)
34+
35+
// Verify plugins are set on pool
36+
if pool.GetPlugins() == nil {
37+
t.Error("plugins should not be nil after SetPlugins")
38+
}
39+
if pool.GetPlugins() != plugins {
40+
t.Error("pool plugins should be the same as the set plugins")
41+
}
42+
43+
// Verify plugins are set on all xclients
44+
for i := 0; i < 3; i++ {
45+
xclient := pool.Get()
46+
if xclient.GetPlugins() == nil {
47+
t.Errorf("xclient %d plugins should not be nil", i)
48+
}
49+
if xclient.GetPlugins() != plugins {
50+
t.Errorf("xclient %d plugins should be the same as the set plugins", i)
51+
}
52+
}
53+
}
54+
55+
func TestXClientPool_GetPlugins(t *testing.T) {
56+
// Create a simple discovery
57+
pairs := []*KVPair{
58+
{Key: "[email protected]:8972", Value: ""},
59+
}
60+
discovery, err := NewMultipleServersDiscovery(pairs)
61+
if err != nil {
62+
t.Fatalf("failed to create discovery: %v", err)
63+
}
64+
defer discovery.Close()
65+
66+
// Create a pool
67+
pool := NewXClientPool(2, "Arith", Failtry, RandomSelect, discovery, DefaultOption)
68+
defer pool.Close()
69+
70+
// Initially, plugins should be nil
71+
if pool.GetPlugins() != nil {
72+
t.Error("plugins should be nil initially")
73+
}
74+
75+
// Create and set plugins
76+
plugins := NewPluginContainer()
77+
tp := &testPlugin{name: "test-plugin"}
78+
plugins.Add(tp)
79+
pool.SetPlugins(plugins)
80+
81+
// Verify GetPlugins returns the correct plugins
82+
retrievedPlugins := pool.GetPlugins()
83+
if retrievedPlugins == nil {
84+
t.Error("plugins should not be nil after SetPlugins")
85+
}
86+
if retrievedPlugins != plugins {
87+
t.Error("GetPlugins should return the same plugins as SetPlugins")
88+
}
89+
90+
// Verify plugins contain the test plugin
91+
allPlugins := retrievedPlugins.All()
92+
if len(allPlugins) != 1 {
93+
t.Errorf("expected 1 plugin, got %d", len(allPlugins))
94+
}
95+
if p, ok := allPlugins[0].(*testPlugin); !ok || p.name != "test-plugin" {
96+
t.Error("plugin should be the test plugin")
97+
}
98+
}
99+
100+
func TestXClientPool_SetPlugins_Concurrent(t *testing.T) {
101+
// Create a simple discovery
102+
pairs := []*KVPair{
103+
{Key: "[email protected]:8972", Value: ""},
104+
}
105+
discovery, err := NewMultipleServersDiscovery(pairs)
106+
if err != nil {
107+
t.Fatalf("failed to create discovery: %v", err)
108+
}
109+
defer discovery.Close()
110+
111+
// Create a pool
112+
pool := NewXClientPool(5, "Arith", Failtry, RandomSelect, discovery, DefaultOption)
113+
defer pool.Close()
114+
115+
// Test concurrent SetPlugins calls
116+
done := make(chan bool, 10)
117+
for i := 0; i < 10; i++ {
118+
go func(id int) {
119+
plugins := NewPluginContainer()
120+
tp := &testPlugin{name: "test-plugin"}
121+
plugins.Add(tp)
122+
pool.SetPlugins(plugins)
123+
done <- true
124+
}(i)
125+
}
126+
127+
// Wait for all goroutines to complete
128+
for i := 0; i < 10; i++ {
129+
<-done
130+
}
131+
132+
// Verify plugins are set
133+
if pool.GetPlugins() == nil {
134+
t.Error("plugins should not be nil after concurrent SetPlugins")
135+
}
136+
}
137+

0 commit comments

Comments
ย (0)