From 44dcf02c697eba456bfeb28cd6587ee5b834d46f Mon Sep 17 00:00:00 2001
From: Artur Melanchyk <artur.melanchyk@gmail.com>
Date: Sat, 23 Mar 2024 19:44:30 +0100
Subject: [PATCH] TSDB: make total lock-free by using atomic

Signed-off-by: Artur Melanchyk <artur.melanchyk@gmail.com>
---
 cmd/promtool/tsdb.go | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/cmd/promtool/tsdb.go b/cmd/promtool/tsdb.go
index 73258754e..b786c9297 100644
--- a/cmd/promtool/tsdb.go
+++ b/cmd/promtool/tsdb.go
@@ -33,6 +33,7 @@ import (
 
 	"github.com/alecthomas/units"
 	"github.com/go-kit/log"
+	"go.uber.org/atomic"
 
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/promql/parser"
@@ -149,8 +150,7 @@ func benchmarkWrite(outPath, samplesFile string, numMetrics, numScrapes int) err
 }
 
 func (b *writeBenchmark) ingestScrapes(lbls []labels.Labels, scrapeCount int) (uint64, error) {
-	var mu sync.Mutex
-	var total uint64
+	var total atomic.Uint64
 
 	for i := 0; i < scrapeCount; i += 100 {
 		var wg sync.WaitGroup
@@ -165,22 +165,21 @@ func (b *writeBenchmark) ingestScrapes(lbls []labels.Labels, scrapeCount int) (u
 
 			wg.Add(1)
 			go func() {
+				defer wg.Done()
+
 				n, err := b.ingestScrapesShard(batch, 100, int64(timeDelta*i))
 				if err != nil {
 					// exitWithError(err)
 					fmt.Println(" err", err)
 				}
-				mu.Lock()
-				total += n
-				mu.Unlock()
-				wg.Done()
+				total.Add(n)
 			}()
 		}
 		wg.Wait()
 	}
 	fmt.Println("ingestion completed")
 
-	return total, nil
+	return total.Load(), nil
 }
 
 func (b *writeBenchmark) ingestScrapesShard(lbls []labels.Labels, scrapeCount int, baset int64) (uint64, error) {