add Close for ActiveQueryTracker to close the file.

Signed-off-by: ianwoolf <btw515wolf2@gmail.com>
This commit is contained in:
ianwoolf 2022-02-27 20:49:33 +08:00 committed by zhangzizhao
parent 7309ac2721
commit 79e4bdee8e
3 changed files with 21 additions and 7 deletions

View File

@ -49,6 +49,7 @@ func TestQueryConcurrency(t *testing.T) {
require.NoError(t, err)
defer os.RemoveAll(dir)
queryTracker := NewActiveQueryTracker(dir, maxConcurrency, nil)
t.Cleanup(queryTracker.Close)
opts := EngineOpts{
Logger: nil,

View File

@ -16,6 +16,7 @@ package promql
import (
"context"
"encoding/json"
"io"
"os"
"path/filepath"
"strings"
@ -31,6 +32,7 @@ type ActiveQueryTracker struct {
mmapedFile []byte
getNextIndex chan int
logger log.Logger
closer io.Closer
maxConcurrent int
}
@ -81,7 +83,7 @@ func logUnfinishedQueries(filename string, filesize int, logger log.Logger) {
}
}
func getMMapedFile(filename string, filesize int, logger log.Logger) ([]byte, error) {
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 {
absPath, pathErr := filepath.Abs(filename)
@ -89,22 +91,22 @@ func getMMapedFile(filename string, filesize int, logger log.Logger) ([]byte, er
absPath = filename
}
level.Error(logger).Log("msg", "Error opening query log file", "file", absPath, "err", err)
return nil, err
return nil, nil, err
}
err = file.Truncate(int64(filesize))
if err != nil {
level.Error(logger).Log("msg", "Error setting filesize.", "filesize", filesize, "err", err)
return nil, err
return nil, nil, err
}
fileAsBytes, err := mmap.Map(file, mmap.RDWR, 0)
if err != nil {
level.Error(logger).Log("msg", "Failed to mmap", "file", filename, "Attempted size", filesize, "err", err)
return nil, err
return nil, nil, err
}
return fileAsBytes, err
return fileAsBytes, file, err
}
func NewActiveQueryTracker(localStoragePath string, maxConcurrent int, logger log.Logger) *ActiveQueryTracker {
@ -116,7 +118,7 @@ func NewActiveQueryTracker(localStoragePath string, maxConcurrent int, logger lo
filename, filesize := filepath.Join(localStoragePath, "queries.active"), 1+maxConcurrent*entrySize
logUnfinishedQueries(filename, filesize, logger)
fileAsBytes, err := getMMapedFile(filename, filesize, logger)
fileAsBytes, closer, err := getMMapedFile(filename, filesize, logger)
if err != nil {
panic("Unable to create mmap-ed active query log")
}
@ -124,6 +126,7 @@ func NewActiveQueryTracker(localStoragePath string, maxConcurrent int, logger lo
copy(fileAsBytes, "[")
activeQueryTracker := ActiveQueryTracker{
mmapedFile: fileAsBytes,
closer: closer,
getNextIndex: make(chan int, maxConcurrent),
logger: logger,
maxConcurrent: maxConcurrent,
@ -198,3 +201,10 @@ func (tracker ActiveQueryTracker) Insert(ctx context.Context, query string) (int
return 0, ctx.Err()
}
}
func (tracker *ActiveQueryTracker) Close() {
if tracker == nil || tracker.closer == nil {
return
}
tracker.closer.Close()
}

View File

@ -110,7 +110,10 @@ func TestMMapFile(t *testing.T) {
filename := file.Name()
defer os.Remove(filename)
fileAsBytes, err := getMMapedFile(filename, 2, nil)
fileAsBytes, closer, err := getMMapedFile(filename, 2, nil)
if err != nil {
t.Cleanup(func() { closer.Close() })
}
require.NoError(t, err)
copy(fileAsBytes, "ab")