Update lib/storage/raw_row.go

Co-authored-by: cubic-dev-ai[bot] <191113872+cubic-dev-ai[bot]@users.noreply.github.com>
Signed-off-by: Artem Fetishev <149964189+rtm0@users.noreply.github.com>

lib/storage/raw_row.go

@@ -236,16 +236,24 @@ func (rrss *rawRowsShards) updateFlushDeadline() {
 	rrss.flushDeadlineMs.Store(time.Now().Add(pendingRowsFlushInterval).UnixMilli())
 }
 
+// flush calls flush callback on all rawRows that are stored in rrss internal
+// buffers and are ready to be flushed.
+//
+// A rawRow is ready to be flushed either when it has spent enough time in rrss
+// internal buffers (see pendingRowsFlushInterval) or if the operation is final
+// and the rawRow needs to be flushed immediately.
+//
+// The flushed rawRows are removed from rrss internal buffers.
 func (rrss *rawRowsShards) flush(flush func(rrs [][]rawRow), isFinal bool) {
-	var dst [][]rawRow
+	rrss.rowssToFlushLock.Lock()
+	defer rrss.rowssToFlushLock.Unlock()
 
+	var dst [][]rawRow
 	currentTimeMs := time.Now().UnixMilli()
 	flushDeadlineMs := rrss.flushDeadlineMs.Load()
 	if isFinal || currentTimeMs >= flushDeadlineMs {
-		rrss.rowssToFlushLock.Lock()
 		dst = rrss.rowssToFlush
 		rrss.rowssToFlush = nil
-		rrss.rowssToFlushLock.Unlock()
 	}
 
 	for i := range rrss.shards {

lib/storage/search_test.go

@@ -8,6 +8,7 @@ import (
 	"reflect"
 	"regexp"
 	"sort"
+	"sync"
 	"testing"
 	"testing/quick"
 	"time"
@@ -191,6 +192,62 @@ func TestSearch_VariousTimeRanges(t *testing.T) {
 	testStorageOpOnVariousTimeRanges(t, f)
 }
 
+func TestStorageAddThenSearchConcurrently(t *testing.T) {
+	defer testRemoveAll(t)
+
+	s := MustOpenStorage(t.Name(), OpenOptions{})
+	defer s.MustClose()
+
+	const numMetrics = 100
+	f := func(workerID int, tr TimeRange) error {
+		mrs := make([]MetricRow, numMetrics)
+		step := (tr.MaxTimestamp - tr.MinTimestamp) / int64(numMetrics)
+		for i := range numMetrics {
+			name := fmt.Sprintf("metric_%04d_%04d", workerID, i)
+			mn := MetricName{MetricGroup: []byte(name)}
+			mrs[i].MetricNameRaw = mn.marshalRaw(nil)
+			mrs[i].Timestamp = tr.MinTimestamp + int64(i)*step
+			mrs[i].Value = float64(i)
+		}
+
+		s.AddRows(mrs, defaultPrecisionBits)
+		s.DebugFlush()
+
+		tfs := NewTagFilters()
+		re := fmt.Sprintf(`metric_%04d.*`, workerID)
+		if err := tfs.Add(nil, []byte(re), false, true); err != nil {
+			return fmt.Errorf("tfs.Add(%q) failed unexpectedly: %w", re, err)
+		}
+		return testAssertSearchResult(s, tr, tfs, mrs)
+	}
+
+	const concurrency = 20
+	var wg sync.WaitGroup
+	errs := make([]error, concurrency)
+	for workerID := range concurrency {
+		wg.Go(func() {
+			for m := time.Month(1); m <= 12; m++ {
+				tr := TimeRange{
+					MinTimestamp: time.Date(2025, m, 1, 0, 0, 0, 0, time.UTC).UnixMilli(),
+					MaxTimestamp: time.Date(2025, m+1, 0, 0, 0, 0, 0, time.UTC).UnixMilli() - 1,
+				}
+				err := f(workerID, tr)
+				if err != nil {
+					errs[workerID] = fmt.Errorf("worker %d failed on tr=%v: %w", workerID, &tr, err)
+					break
+				}
+			}
+		})
+	}
+	wg.Wait()
+
+	for _, err := range errs {
+		if err != nil {
+			t.Errorf("%s", err)
+		}
+	}
+}
+
 func testSearchInternal(s *Storage, tr TimeRange, mrs []MetricRow) error {
 	for i := range 10 {
 		// Prepare TagFilters for search.