From 519f6be5429217a1db4bca2cb0821cc330b0cca8 Mon Sep 17 00:00:00 2001
From: yomek33 <i3landme@gmail.com>
Date: Tue, 16 Jul 2024 05:58:23 +0900
Subject: [PATCH] Add series gradual increase and decrease functionlity

Signed-off-by: yomek33 <i3landme@gmail.com>
---
 cmd/avalanche.go      |  2 +-
 metrics/serve.go      | 22 +++++++++++++++
 metrics/serve_test.go | 62 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 85 insertions(+), 1 deletion(-)

diff --git a/cmd/avalanche.go b/cmd/avalanche.go
index 6bcd996..5eb000e 100644
--- a/cmd/avalanche.go
+++ b/cmd/avalanche.go
@@ -41,7 +41,7 @@ var (
 	labelInterval        = kingpin.Flag("series-interval", "Change series_id label values every {interval} seconds.").Default("60").Int()
 	metricInterval       = kingpin.Flag("metric-interval", "Change __name__ label values every {interval} seconds.").Default("120").Int()
 	seriesChangeInterval = kingpin.Flag("series-change-interval", "Change the number of series every {interval} seconds.").Default("10").Int()
-	operationMode        = kingpin.Flag("operation-mode", "Mode of operation: 'gradual-change', 'series-spike', 'double-halve'").Default("gradual-change").String()
+	operationMode        = kingpin.Flag("operation-mode", "Mode of operation: 'gradual-change', 'series-spike', 'double-halve'").Default("default").String()
 	port                 = kingpin.Flag("port", "Port to serve at").Default("9001").Int()
 	remoteURL            = kingpin.Flag("remote-url", "URL to send samples via remote_write API.").URL()
 	remotePprofURLs      = kingpin.Flag("remote-pprof-urls", "a list of urls to download pprofs during the remote write: --remote-pprof-urls=http://127.0.0.1:10902/debug/pprof/heap --remote-pprof-urls=http://127.0.0.1:10902/debug/pprof/profile").URLList()
diff --git a/metrics/serve.go b/metrics/serve.go
index 6895d2e..c090a08 100644
--- a/metrics/serve.go
+++ b/metrics/serve.go
@@ -184,6 +184,28 @@ func RunMetrics(metricCount, labelCount, seriesCount, seriesChangeRate, metricLe
 				}
 			}
 		}()
+
+	case "gradual-change":
+		go func() {
+			for tick := range changeSeriesTick.C {
+				fmt.Printf("%v: Adjusting series count. New count: %d\n", tick, seriesCount)
+				metricsMux.Lock()
+				seriesCount += seriesChangeRate
+				if seriesCount < 1 {
+					seriesCount = 1
+				}
+				unregisterMetrics()
+				registerMetrics(metricCount, metricLength, metricCycle, labelKeys)
+				cycleValues(labelKeys, labelValues, seriesCount, seriesCycle)
+				metricsMux.Unlock()
+
+				select {
+				case updateNotify <- struct{}{}:
+				default:
+				}
+			}
+		}()
+	default:
 	}
 
 	go func() {
diff --git a/metrics/serve_test.go b/metrics/serve_test.go
index f63aef4..988ccd6 100644
--- a/metrics/serve_test.go
+++ b/metrics/serve_test.go
@@ -75,3 +75,65 @@ func TestRunMetricsSeriesCountChangeDoubleHalve(t *testing.T) {
 		t.Fatal("Did not receive update notification for series count halving in time")
 	}
 }
+
+func TestRunMetricsGradualChange(t *testing.T) {
+	const (
+		initialSeriesCount    = 30
+		metricCount           = 1
+		labelCount            = 1
+		seriesChangeRate      = -10
+		metricLength          = 1
+		labelLength           = 1
+		valueInterval         = 1
+		seriesInterval        = 1
+		metricInterval        = 1
+		seriesChangeInterval  = 3
+		operationMode         = "gradual-change"
+		constLabel            = "constLabel=test"
+		updateNotifyTimeout   = 4 * time.Second
+		waitTimeBetweenChecks = 3 * time.Second
+	)
+
+	stop := make(chan struct{})
+	defer close(stop)
+
+	promRegistry = prometheus.NewRegistry()
+
+	updateNotify, err := RunMetrics(metricCount, labelCount, initialSeriesCount, seriesChangeRate, metricLength, labelLength, valueInterval, seriesInterval, metricInterval, seriesChangeInterval, operationMode, []string{constLabel}, stop)
+	assert.NoError(t, err)
+
+	initialCount := countSeries(t, promRegistry)
+	expectedInitialCount := initialSeriesCount
+	assert.Equal(t, expectedInitialCount, initialCount, "Initial series count should be %d but got %d", expectedInitialCount, initialCount)
+
+	select {
+	case <-updateNotify:
+		time.Sleep(waitTimeBetweenChecks)
+		updatedCount := countSeries(t, promRegistry)
+		expectedCount := initialSeriesCount + seriesChangeRate
+		assert.Equal(t, expectedCount, updatedCount, "1 Decreased series count should be %d but got %d", expectedCount, updatedCount)
+	case <-time.After(updateNotifyTimeout):
+		t.Fatal("Did not receive update notification for series count doubling in time")
+	}
+
+	select {
+	case <-updateNotify:
+		time.Sleep(waitTimeBetweenChecks)
+		updatedCount := countSeries(t, promRegistry)
+		expectedCount := initialSeriesCount + seriesChangeRate*2
+		assert.Equal(t, expectedCount, updatedCount, "2 Decreased series count should be %d but got %d", expectedCount, updatedCount)
+	case <-time.After(updateNotifyTimeout):
+		t.Fatal("Did not receive update notification for series count doubling in time")
+	}
+
+	// Test for the minimum value of the series is 1
+	select {
+	case <-updateNotify:
+		time.Sleep(waitTimeBetweenChecks)
+		updatedCount := countSeries(t, promRegistry)
+		expectedCount := 1
+		assert.Equal(t, expectedCount, updatedCount, "3 Decreased series count should be %d but got %d", expectedCount, updatedCount)
+	case <-time.After(updateNotifyTimeout):
+		t.Fatal("Did not receive update notification for series count doubling in time")
+	}
+}