diff --git a/discovery/generator.go b/discovery/generator.go new file mode 100644 index 0000000000..d95c5ce3d8 --- /dev/null +++ b/discovery/generator.go @@ -0,0 +1,37 @@ +package discovery + +import ( + "fmt" + "regexp" + "strconv" +) + +// +// IP generator +// +func Generate(pattern string) []string { + re, _ := regexp.Compile(`\[(.+):(.+)\]`) + submatch := re.FindStringSubmatch(pattern) + if submatch == nil { + return []string{pattern} + } + + from, err := strconv.Atoi(submatch[1]) + if err != nil { + return []string{pattern} + } + to, err := strconv.Atoi(submatch[2]) + if err != nil { + return []string{pattern} + } + + template := re.ReplaceAllString(pattern, "%d") + + result := make([]string, 0) + for val := from; val <= to; val++ { + entry := fmt.Sprintf(template, val) + result = append(result, entry) + } + + return result +} diff --git a/discovery/generator_test.go b/discovery/generator_test.go new file mode 100644 index 0000000000..827d1ef55d --- /dev/null +++ b/discovery/generator_test.go @@ -0,0 +1,41 @@ +package discovery + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestGeneratorNotGenerate(t *testing.T) { + ips := Generate("127.0.0.1") + assert.Equal(t, len(ips), 1) + assert.Equal(t, ips[0], "127.0.0.1") +} + +func TestGeneratorWithPortNotGenerate(t *testing.T) { + ips := Generate("127.0.0.1:8080") + assert.Equal(t, len(ips), 1) + assert.Equal(t, ips[0], "127.0.0.1:8080") +} + +func TestGeneratorMatchFailedNotGenerate(t *testing.T) { + ips := Generate("127.0.0.[1]") + assert.Equal(t, len(ips), 1) + assert.Equal(t, ips[0], "127.0.0.[1]") +} + +func TestGeneratorWithPort(t *testing.T) { + ips := Generate("127.0.0.[1:11]:2375") + assert.Equal(t, len(ips), 11) + assert.Equal(t, ips[0], "127.0.0.1:2375") + assert.Equal(t, ips[1], "127.0.0.2:2375") + assert.Equal(t, ips[2], "127.0.0.3:2375") + assert.Equal(t, ips[3], "127.0.0.4:2375") + assert.Equal(t, ips[4], "127.0.0.5:2375") + assert.Equal(t, ips[5], "127.0.0.6:2375") + assert.Equal(t, ips[6], "127.0.0.7:2375") + assert.Equal(t, ips[7], "127.0.0.8:2375") + assert.Equal(t, ips[8], "127.0.0.9:2375") + assert.Equal(t, ips[9], "127.0.0.10:2375") + assert.Equal(t, ips[10], "127.0.0.11:2375") +} diff --git a/discovery/nodes/nodes.go b/discovery/nodes/nodes.go index 099f00494e..b3bde4f6f4 100644 --- a/discovery/nodes/nodes.go +++ b/discovery/nodes/nodes.go @@ -15,12 +15,14 @@ func init() { } func (s *NodesDiscoveryService) Initialize(uris string, _ int) error { - for _, ip := range strings.Split(uris, ",") { - entry, err := discovery.NewEntry(ip) - if err != nil { - return err + for _, input := range strings.Split(uris, ",") { + for _, ip := range discovery.Generate(input) { + entry, err := discovery.NewEntry(ip) + if err != nil { + return err + } + s.entries = append(s.entries, entry) } - s.entries = append(s.entries, entry) } return nil diff --git a/discovery/nodes/nodes_test.go b/discovery/nodes/nodes_test.go index f89ccd145e..96a3cf7fd2 100644 --- a/discovery/nodes/nodes_test.go +++ b/discovery/nodes/nodes_test.go @@ -14,6 +14,17 @@ func TestInitialise(t *testing.T) { assert.Equal(t, discovery.entries[1].String(), "2.2.2.2:2222") } +func TestInitialiseWithPattern(t *testing.T) { + discovery := &NodesDiscoveryService{} + discovery.Initialize("1.1.1.[1:2]:1111,2.2.2.[2:4]:2222", 0) + assert.Equal(t, len(discovery.entries), 5) + assert.Equal(t, discovery.entries[0].String(), "1.1.1.1:1111") + assert.Equal(t, discovery.entries[1].String(), "1.1.1.2:1111") + assert.Equal(t, discovery.entries[2].String(), "2.2.2.2:2222") + assert.Equal(t, discovery.entries[3].String(), "2.2.2.3:2222") + assert.Equal(t, discovery.entries[4].String(), "2.2.2.4:2222") +} + func TestRegister(t *testing.T) { discovery := &NodesDiscoveryService{} assert.Error(t, discovery.Register("0.0.0.0"))