From b7f2f3c3ac90f2347de6112c185a4e470e7ae8a6 Mon Sep 17 00:00:00 2001
From: Oleg Zaytsev <mail@olegzaytsev.com>
Date: Tue, 30 Jul 2024 10:19:56 +0200
Subject: [PATCH] Add BenchmarkLoadRealWLs

This benchmark runs on real WLs rather than fake generated ones.

Signed-off-by: Oleg Zaytsev <mail@olegzaytsev.com>
---
 tsdb/head_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/tsdb/head_test.go b/tsdb/head_test.go
index c192c8a078..09927c23c6 100644
--- a/tsdb/head_test.go
+++ b/tsdb/head_test.go
@@ -23,6 +23,7 @@ import (
 	"path"
 	"path/filepath"
 	"reflect"
+	"runtime/pprof"
 	"sort"
 	"strconv"
 	"strings"
@@ -89,6 +90,43 @@ func newTestHeadWithOptions(t testing.TB, compressWAL wlog.CompressionType, opts
 	return h, wal
 }
 
+// BenchmarkLoadRealWLs will be skipped unless the BENCHMARK_LOAD_REAL_WLS_DIR environment variable is set.
+// BENCHMARK_LOAD_REAL_WLS_DIR should be the folder where `wal` and `chunks_head` are located.
+// Optionally, BENCHMARK_LOAD_REAL_WLS_PROFILE can be set to a file path to write a CPU profile.
+func BenchmarkLoadRealWLs(b *testing.B) {
+	dir := os.Getenv("BENCHMARK_LOAD_REAL_WLS_DIR")
+	if dir == "" {
+		b.Skipped()
+	}
+
+	profileFile := os.Getenv("BENCHMARK_LOAD_REAL_WLS_PROFILE")
+	if profileFile != "" {
+		b.Logf("Will profile in %s", profileFile)
+		f, err := os.Create(profileFile)
+		require.NoError(b, err)
+		b.Cleanup(func() { f.Close() })
+		require.NoError(b, pprof.StartCPUProfile(f))
+		b.Cleanup(pprof.StopCPUProfile)
+	}
+
+	wal, err := wlog.New(nil, nil, filepath.Join(dir, "wal"), wlog.CompressionNone)
+	require.NoError(b, err)
+	b.Cleanup(func() { wal.Close() })
+
+	wbl, err := wlog.New(nil, nil, filepath.Join(dir, "wbl"), wlog.CompressionNone)
+	require.NoError(b, err)
+	b.Cleanup(func() { wbl.Close() })
+
+	// Load the WAL.
+	for i := 0; i < b.N; i++ {
+		opts := DefaultHeadOptions()
+		opts.ChunkDirRoot = dir
+		h, err := NewHead(nil, nil, wal, wbl, opts, nil)
+		require.NoError(b, err)
+		h.Init(0)
+	}
+}
+
 func BenchmarkCreateSeries(b *testing.B) {
 	series := genSeries(b.N, 10, 0, 0)
 	h, _ := newTestHead(b, 10000, wlog.CompressionNone, false)