From 03873a779e634dddcf3a85b3e5fecfaa4aaab779 Mon Sep 17 00:00:00 2001
From: Alexei Shevchenko <al3x.s0rg@gmail.com>
Date: Sun, 3 Sep 2023 22:51:41 +0300
Subject: [PATCH] weight for clusterization rules (#29)

---
 README.md                      | 17 ++++++++++-------
 examples/cluster.json          |  2 +-
 internal/graph/cluster.go      | 18 +++++++++++++++---
 internal/graph/cluster_test.go | 28 ++++++++++++++++++++++++++++
 4 files changed, 54 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index 0e61449..f8cc442 100644
--- a/README.md
+++ b/README.md
@@ -119,19 +119,22 @@ See [csv2meta.py](examples/csv2meta.py) for example how to create such `json` fo
 
 # clusterization rules
 
-You can join your services into `clusters` by exposed ports, with clusterization rules, for example:
+You can join your services into `clusters` by exposed ports, in `dot` or `structurizr` output formats.
+With clusterization rules, in `json`:
 
 ```json
 [
-    {"name": "ingress", "ports": ["80/tcp", "443/tcp"]},
-    {"name": "backend", "ports": ["8080-8090/tcp"]},
-    {"name": "store", "ports": ["3306/tcp", "5432/tcp"]},
-    {"name": "redis", "ports": ["6379/tcp"]},
-    {"name": "queue", "ports": ["9092/tcp", "4222/tcp"]}
+    {
+        "name": "rule-name",
+        "weight": 1,
+        "ports": ["port/kind", "ports-range/kind"]
+    },
+    ...
 ]
 ```
 
-Source: [cluster.json](examples/cluster.json)
+Weight can be omitted, if not specified it equals 1.
+See: [cluster.json](examples/cluster.json) for detailed example.
 
 # features
 
diff --git a/examples/cluster.json b/examples/cluster.json
index 5d05f17..2ce6213 100644
--- a/examples/cluster.json
+++ b/examples/cluster.json
@@ -1,6 +1,6 @@
 [
     {"name": "ingress", "ports": ["80/tcp", "443/tcp"]},
-    {"name": "backend", "ports": ["8080-8090/tcp"]},
+    {"name": "backend", "weight": 2, "ports": ["8080-8090/tcp"]},
     {"name": "store", "ports": ["3306/tcp", "5432/tcp"]},
     {"name": "redis", "ports": ["6379/tcp"]},
     {"name": "queue", "ports": ["9092/tcp", "4222/tcp"]}
diff --git a/internal/graph/cluster.go b/internal/graph/cluster.go
index e8ad9ac..33d627b 100644
--- a/internal/graph/cluster.go
+++ b/internal/graph/cluster.go
@@ -21,12 +21,14 @@ var (
 
 type (
 	ruleJSON struct {
-		Name  string   `json:"name"`
-		Ports []string `json:"ports"`
+		Name   string   `json:"name"`
+		Ports  []string `json:"ports"`
+		Weight int      `json:"weight"`
 	}
 
 	ClusterBuilder struct {
 		builder NamedBuilderWriter
+		weights map[string]int
 		nodes   map[string]*node.Node
 		index   map[string]map[int]string
 		cluster map[string]map[string]node.Ports
@@ -36,6 +38,7 @@ type (
 func NewClusterBuilder(b NamedBuilderWriter) *ClusterBuilder {
 	return &ClusterBuilder{
 		builder: b,
+		weights: make(map[string]int),
 		nodes:   make(map[string]*node.Node),
 		index:   make(map[string]map[int]string),
 		cluster: make(map[string]map[string]node.Ports),
@@ -114,6 +117,13 @@ func (cb *ClusterBuilder) FromReader(r io.Reader) (err error) {
 	for i := 0; i < len(rules); i++ {
 		rule := &rules[i]
 
+		weight := rule.Weight
+		if weight == 0 {
+			weight = 1
+		}
+
+		cb.weights[rule.Name] = weight
+
 		for j := 0; j < len(rule.Ports); j++ {
 			proto, ports, perr := parseRulePorts(rule.Ports[j])
 			if perr != nil {
@@ -165,7 +175,9 @@ func (cb *ClusterBuilder) Match(n *node.Node) (cluster string, ok bool) {
 	smatches := make([]*match, 0, len(matches))
 
 	for k, n := range matches {
-		smatches = append(smatches, &match{Name: k, Weight: n})
+		w := cb.weights[k]
+
+		smatches = append(smatches, &match{Name: k, Weight: n * w})
 	}
 
 	switch len(smatches) {
diff --git a/internal/graph/cluster_test.go b/internal/graph/cluster_test.go
index 1a1fd6a..a547651 100644
--- a/internal/graph/cluster_test.go
+++ b/internal/graph/cluster_test.go
@@ -135,6 +135,34 @@ func TestClusterMatch(t *testing.T) {
 	}
 }
 
+func TestClusterMatchWeight(t *testing.T) {
+	t.Parallel()
+
+	const clusterRulesWeight = `[{"name": "foo", "ports": ["80-80/tcp"]},
+{"name": "bar", "weight": 2, "ports": ["22/tcp", "443/tcp"]}]`
+
+	testNode := &node.Node{Ports: []node.Port{
+		{Kind: "tcp", Value: 22},
+		{Kind: "tcp", Value: 80},
+		{Kind: "tcp", Value: 8080},
+	}}
+
+	ca := graph.NewClusterBuilder(nil)
+
+	if err := ca.FromReader(bytes.NewBufferString(clusterRulesWeight)); err != nil {
+		t.Fatal(err)
+	}
+
+	m, ok := ca.Match(testNode)
+	if !ok {
+		t.Fail()
+	}
+
+	if m != "bar" {
+		t.Fail()
+	}
+}
+
 func TestClusterBuilder(t *testing.T) {
 	t.Parallel()