From aabe4d6e4ab39d4d87611668312ec4d93616b5b8 Mon Sep 17 00:00:00 2001
From: Arve Knudsen <arve.knudsen@gmail.com>
Date: Mon, 29 Apr 2024 16:16:51 +0200
Subject: [PATCH] promql.ActiveQueryTracker: Unmap mmapped file when done

Signed-off-by: Arve Knudsen <arve.knudsen@gmail.com>
---
 promql/engine_test.go       |  4 +++-
 promql/query_logger.go      | 32 ++++++++++++++++++++++++++++----
 promql/query_logger_test.go | 27 ++++++++++++++-------------
 3 files changed, 45 insertions(+), 18 deletions(-)

diff --git a/promql/engine_test.go b/promql/engine_test.go
index cc9185546..485239399 100644
--- a/promql/engine_test.go
+++ b/promql/engine_test.go
@@ -59,7 +59,9 @@ func TestQueryConcurrency(t *testing.T) {
 	require.NoError(t, err)
 	defer os.RemoveAll(dir)
 	queryTracker := promql.NewActiveQueryTracker(dir, maxConcurrency, nil)
-	t.Cleanup(queryTracker.Close)
+	t.Cleanup(func() {
+		require.NoError(t, queryTracker.Close())
+	})
 
 	opts := promql.EngineOpts{
 		Logger:             nil,
diff --git a/promql/query_logger.go b/promql/query_logger.go
index 7ddd8c2d5..76528f958 100644
--- a/promql/query_logger.go
+++ b/promql/query_logger.go
@@ -16,6 +16,7 @@ package promql
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"io"
 	"os"
 	"path/filepath"
@@ -36,6 +37,8 @@ type ActiveQueryTracker struct {
 	maxConcurrent int
 }
 
+var _ io.Closer = &ActiveQueryTracker{}
+
 type Entry struct {
 	Query     string `json:"query"`
 	Timestamp int64  `json:"timestamp_sec"`
@@ -83,6 +86,23 @@ func logUnfinishedQueries(filename string, filesize int, logger log.Logger) {
 	}
 }
 
+type mmapedFile struct {
+	f io.Closer
+	m mmap.MMap
+}
+
+func (f *mmapedFile) Close() error {
+	err := f.m.Unmap()
+	if fErr := f.f.Close(); fErr != nil && err == nil {
+		return fmt.Errorf("close mmapedFile.f: %w", fErr)
+	}
+
+	if err != nil {
+		return fmt.Errorf("mmapedFile: unmapping: %w", err)
+	}
+	return nil
+}
+
 func getMMapedFile(filename string, filesize int, logger log.Logger) ([]byte, io.Closer, error) {
 	file, err := os.OpenFile(filename, os.O_CREATE|os.O_RDWR|os.O_TRUNC, 0o666)
 	if err != nil {
@@ -108,7 +128,7 @@ func getMMapedFile(filename string, filesize int, logger log.Logger) ([]byte, io
 		return nil, nil, err
 	}
 
-	return fileAsBytes, file, err
+	return fileAsBytes, &mmapedFile{f: file, m: fileAsBytes}, err
 }
 
 func NewActiveQueryTracker(localStoragePath string, maxConcurrent int, logger log.Logger) *ActiveQueryTracker {
@@ -204,9 +224,13 @@ func (tracker ActiveQueryTracker) Insert(ctx context.Context, query string) (int
 	}
 }
 
-func (tracker *ActiveQueryTracker) Close() {
+// Close closes tracker.
+func (tracker *ActiveQueryTracker) Close() error {
 	if tracker == nil || tracker.closer == nil {
-		return
+		return nil
 	}
-	tracker.closer.Close()
+	if err := tracker.closer.Close(); err != nil {
+		return fmt.Errorf("close ActiveQueryTracker.closer: %w", err)
+	}
+	return nil
 }
diff --git a/promql/query_logger_test.go b/promql/query_logger_test.go
index 376d61b64..7bd93781e 100644
--- a/promql/query_logger_test.go
+++ b/promql/query_logger_test.go
@@ -16,6 +16,7 @@ package promql
 import (
 	"context"
 	"os"
+	"path/filepath"
 	"testing"
 
 	"github.com/grafana/regexp"
@@ -104,26 +105,26 @@ func TestIndexReuse(t *testing.T) {
 }
 
 func TestMMapFile(t *testing.T) {
-	file, err := os.CreateTemp("", "mmapedFile")
+	dir := t.TempDir()
+	fpath := filepath.Join(dir, "mmapedFile")
+	const data = "ab"
+
+	fileAsBytes, closer, err := getMMapedFile(fpath, 2, nil)
 	require.NoError(t, err)
+	copy(fileAsBytes, data)
+	require.NoError(t, closer.Close())
 
-	filename := file.Name()
-	defer os.Remove(filename)
-
-	fileAsBytes, _, err := getMMapedFile(filename, 2, nil)
-
-	require.NoError(t, err)
-	copy(fileAsBytes, "ab")
-
-	f, err := os.Open(filename)
+	f, err := os.Open(fpath)
 	require.NoError(t, err)
+	t.Cleanup(func() {
+		_ = f.Close()
+	})
 
 	bytes := make([]byte, 4)
 	n, err := f.Read(bytes)
-	require.Equal(t, 2, n)
 	require.NoError(t, err, "Unexpected error while reading file.")
-
-	require.Equal(t, fileAsBytes, bytes[:2], "Mmap failed")
+	require.Equal(t, 2, n)
+	require.Equal(t, []byte(data), bytes[:2], "Mmap failed")
 }
 
 func TestParseBrokenJSON(t *testing.T) {