Update lib/persistentqueue/persistentqueue.go

Co-authored-by: cubic-dev-ai[bot] <191113872+cubic-dev-ai[bot]@users.noreply.github.com>
Signed-off-by: Nikolay <nik@victoriametrics.com>

docs/victoriametrics/changelog/CHANGELOG.md

@@ -28,6 +28,7 @@ See also [LTS releases](https://docs.victoriametrics.com/victoriametrics/lts-rel
 
 * FEATURE: [MetricsQL](https://docs.victoriametrics.com/victoriametrics/metricsql/): support `fill` modifiers to allow missing series on either side of a binary operation to be filled with a provided default value. See [#10598](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/10598).
 * FEATURE: [vmagent](https://docs.victoriametrics.com/victoriametrics/vmagent/) and [vmsingle](https://docs.victoriametrics.com/victoriametrics/single-server-victoriametrics/): Improve background discovery performance for [http_sd](https://docs.victoriametrics.com/victoriametrics/sd_configs/#http_sd_configs) discovery. See [#8838](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/8838).
+* FEATURE: [vmagent](https://docs.victoriametrics.com/victoriametrics/vmagent/): restore broken persistent queue chunk file from the last valid written block in case of ungraceful vmagent shutdown. See [#11192](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/11192).
 * FEATURE: [vmagent](https://docs.victoriametrics.com/victoriametrics/vmagent/) and [vmsingle](https://docs.victoriametrics.com/victoriametrics/single-server-victoriametrics/): allow overriding `max_scrape_size` on a per-target basis via the `__max_scrape_size__` label during target relabeling. See [#11188](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/11188).
 * FEATURE: [vmstorage](https://docs.victoriametrics.com/victoriametrics/cluster-victoriametrics/) and [vmsingle](https://docs.victoriametrics.com/victoriametrics/single-server-victoriametrics/): add `-maxBackfillAge` command-line flag for limiting ingestion of samples with historical timestamps, for example, when older data has been moved between storage tiers (nvme/hdd, hot/cold). See [#11199](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/11199). Thanks to @AshwinRamaniPsg for contribution.
 

lib/persistentqueue/persistentqueue.go

@@ -25,6 +25,8 @@ const MaxBlockSize = 32 * 1024 * 1024
 // DefaultChunkFileSize represents default chunk file size
 const DefaultChunkFileSize = (MaxBlockSize + 8) * 16
 
+const blockHeaderSize = 8
+
 var chunkFileNameRegex = regexp.MustCompile("^[0-9A-F]{16}$")
 
 // queue represents persistent queue.
@@ -129,7 +131,7 @@ func mustOpen(path, name string, maxPendingBytes int64) *queue {
 }
 
 func mustOpenInternal(path, name string, chunkFileSize, maxBlockSize, maxPendingBytes uint64) *queue {
-	if chunkFileSize < 8 || chunkFileSize-8 < maxBlockSize {
+	if chunkFileSize < blockHeaderSize || chunkFileSize-blockHeaderSize < maxBlockSize {
 		logger.Panicf("BUG: too small chunkFileSize=%d for maxBlockSize=%d; chunkFileSize must fit at least one block", chunkFileSize, maxBlockSize)
 	}
 	if maxBlockSize <= 0 {
@@ -280,14 +282,30 @@ func tryOpeningQueue(path, name string, chunkFileSize, maxBlockSize, maxPendingB
 			q.writerFlushedOffset = mi.WriterOffset
 			if fileSize := fs.MustFileSize(q.writerPath); fileSize != q.writerLocalOffset {
 				if fileSize < q.writerLocalOffset {
-					logger.Errorf("%q size (%d bytes) is smaller than the writer offset (%d bytes); removing the file",
-						q.writerPath, fileSize, q.writerLocalOffset)
-					fs.MustRemovePath(q.writerPath)
-					continue
+					validChunkFileSize := mustGetChunkValidDataSize(q.writerPath, q.maxBlockSize)
+					logger.Warnf("%q size (%d bytes) is smaller than the writer offset (%d bytes); "+
+						"this may be the case on unclean shutdown (OOM, `kill -9`, hardware reset); resetting writer to fileSize: %d",
+						q.writerPath, fileSize, q.writerLocalOffset, validChunkFileSize)
+					mi.WriterOffset = offset + validChunkFileSize
+					q.writerOffset = mi.WriterOffset
+					q.writerLocalOffset = mi.WriterOffset % q.chunkFileSize
+					q.writerFlushedOffset = mi.WriterOffset
+					// The recovered writer offset may end up smaller than mi.ReaderOffset.
+					// Do not clamp the reader offset in this case: it means the reader has already
+					// consumed data that is now lost, so the remaining queue state cannot be trusted.
+					// The readerOffset > writerOffset check below will fail queue opening,
+					// and the caller (mustOpen) will drop the queue and recreate it from scratch.
+				} else {
+					logger.Warnf("%q size (%d bytes) is bigger than writer offset (%d bytes); "+
+						"this may be the case on unclean shutdown (OOM, `kill -9`, hardware reset); trying to fix it by adjusting fileSize to %d",
+						q.writerPath, fileSize, q.writerLocalOffset, q.writerLocalOffset)
+				}
+				if err := os.Truncate(q.writerPath, int64(q.writerLocalOffset)); err != nil {
+					logger.Panicf("FATAL: cannot truncate chunk: %q to size: %d: %s", q.writerPath, q.writerLocalOffset, err)
+				}
+				if err := mi.WriteToFile(metainfoPath); err != nil {
+					logger.Panicf("FATAL: cannot update metainfo file: %q: %s", metainfoPath, err)
 				}
-				logger.Warnf("%q size (%d bytes) is bigger than writer offset (%d bytes); "+
-					"this may be the case on unclean shutdown (OOM, `kill -9`, hardware reset); trying to fix it by adjusting fileSize to %d",
-					q.writerPath, fileSize, q.writerLocalOffset, q.writerLocalOffset)
 			}
 			w, err := filestream.OpenWriterAt(q.writerPath, int64(q.writerLocalOffset), false)
 			if err != nil {
@@ -357,7 +375,7 @@ func (q *queue) MustWriteBlock(block []byte) {
 	}
 	if q.maxPendingBytes > 0 {
 		// Drain the oldest blocks until the number of pending bytes becomes enough for the block.
-		blockSize := uint64(len(block) + 8)
+		blockSize := uint64(len(block) + blockHeaderSize)
 		maxPendingBytes := q.maxPendingBytes
 		if blockSize < maxPendingBytes {
 			maxPendingBytes -= blockSize
@@ -395,7 +413,7 @@ func (q *queue) writeBlock(block []byte) error {
 	defer func() {
 		writeDurationSeconds.Add(time.Since(startTime).Seconds())
 	}()
-	if q.writerLocalOffset+q.maxBlockSize+8 > q.chunkFileSize {
+	if q.writerLocalOffset+q.maxBlockSize+blockHeaderSize > q.chunkFileSize {
 		if err := q.nextChunkFileForWrite(); err != nil {
 			return fmt.Errorf("cannot create next chunk file: %w", err)
 		}
@@ -408,7 +426,7 @@ func (q *queue) writeBlock(block []byte) error {
 	err := q.write(header.B)
 	headerBufPool.Put(header)
 	if err != nil {
-		return fmt.Errorf("cannot write header with size 8 bytes to %q: %w", q.writerPath, err)
+		return fmt.Errorf("cannot write header with size %d bytes to %q: %w", blockHeaderSize, q.writerPath, err)
 	}
 
 	// Write block contents.
@@ -469,7 +487,7 @@ func (q *queue) readBlock(dst []byte) ([]byte, error) {
 	defer func() {
 		readDurationSeconds.Add(time.Since(startTime).Seconds())
 	}()
-	if q.readerLocalOffset+q.maxBlockSize+8 > q.chunkFileSize {
+	if q.readerLocalOffset+q.maxBlockSize+blockHeaderSize > q.chunkFileSize {
 		if err := q.nextChunkFileForRead(); err != nil {
 			return dst, fmt.Errorf("cannot open next chunk file: %w", err)
 		}
@@ -478,12 +496,12 @@ func (q *queue) readBlock(dst []byte) ([]byte, error) {
 again:
 	// Read block len.
 	header := headerBufPool.Get()
-	header.B = bytesutil.ResizeNoCopyMayOverallocate(header.B, 8)
+	header.B = bytesutil.ResizeNoCopyMayOverallocate(header.B, blockHeaderSize)
 	err := q.readFull(header.B)
 	blockLen := encoding.UnmarshalUint64(header.B)
 	headerBufPool.Put(header)
 	if err != nil {
-		logger.Errorf("skipping corrupted %q, since header with size 8 bytes cannot be read from it: %s", q.readerPath, err)
+		logger.Errorf("skipping corrupted %q, since header with size %d bytes cannot be read from it: %s", q.readerPath, blockHeaderSize, err)
 		if err := q.skipBrokenChunkFile(); err != nil {
 			return dst, err
 		}
@@ -669,3 +687,39 @@ func (mi *metainfo) ReadFromFile(path string) error {
 	}
 	return nil
 }
+
+func mustGetChunkValidDataSize(filePath string, maxBlockSize uint64) uint64 {
+	f, err := os.Open(filePath)
+	if err != nil {
+		logger.Panicf("FATAL: cannot open file: %s", err)
+	}
+	defer fs.MustClose(f)
+
+	fi, err := f.Stat()
+	if err != nil {
+		logger.Panicf("FATAL: cannot read file stat: %s", err)
+	}
+	fileSize := uint64(fi.Size())
+
+	var offset uint64
+	var header [blockHeaderSize]byte
+	for {
+		if offset+blockHeaderSize > fileSize {
+			return offset
+		}
+		if _, err := io.ReadFull(f, header[:]); err != nil {
+			return offset
+		}
+		blockLen := encoding.UnmarshalUint64(header[:])
+		if blockLen == 0 || blockLen > maxBlockSize {
+			return offset
+		}
+		if offset+blockHeaderSize+blockLen > fileSize {
+			return offset
+		}
+		offset += blockHeaderSize + blockLen
+		if _, err := f.Seek(int64(offset), io.SeekStart); err != nil {
+			return offset
+		}
+	}
+}

lib/persistentqueue/persistentqueue_test.go

@@ -137,6 +137,243 @@ func TestQueueOpen(t *testing.T) {
 		q.MustClose()
 		fs.MustRemoveDir(path)
 	})
+	t.Run("damaged-writer-file-tail", func(t *testing.T) {
+		path := "damaged-writer-file-tail"
+		fs.MustRemoveDir(path)
+
+		q := mustOpen(path, "foobar", 0)
+		block1 := []byte("valid block 1")
+		block2 := []byte("valid block 2")
+		q.MustWriteBlock(block1)
+		q.MustWriteBlock(block2)
+		q.MustClose()
+
+		chunkPath := filepath.Join(path, fmt.Sprintf("%016X", 0))
+		chunkFileSize := fs.MustFileSize(chunkPath)
+
+		if err := os.Truncate(chunkPath, int64(chunkFileSize)-1); err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+
+		q = mustOpen(path, "foobar", 0)
+
+		var buf []byte
+		var ok bool
+		buf, ok = q.MustReadBlockNonblocking(buf[:0])
+		if !ok {
+			t.Fatalf("expected block to be readable")
+		}
+		if string(buf) != string(block1) {
+			t.Fatalf("unexpected block got: %q, want: %q", buf, block1)
+		}
+
+		_, ok = q.MustReadBlockNonblocking(buf[:0])
+		if ok {
+			t.Fatalf("expected second block to be dropped")
+		}
+		q.MustClose()
+
+		expectedChunkFileSize := uint64(blockHeaderSize + len(block1))
+		chunkFileSize = fs.MustFileSize(chunkPath)
+		if chunkFileSize != expectedChunkFileSize {
+			t.Fatalf("unexpected chunk file size: got %d; want %d", chunkFileSize, expectedChunkFileSize)
+		}
+
+		fs.MustRemoveDir(path)
+	})
+	t.Run("damaged-writer-file-extra-tail", func(t *testing.T) {
+		path := "damaged-writer-file-extra-tail"
+		fs.MustRemoveDir(path)
+
+		q := mustOpen(path, "foobar", 0)
+		blocks := [][]byte{
+			[]byte("valid block 1"),
+			[]byte("valid block 2"),
+		}
+		for _, block := range blocks {
+			q.MustWriteBlock(block)
+		}
+		q.MustClose()
+
+		chunkPath := filepath.Join(path, fmt.Sprintf("%016X", 0))
+		originChunkFileSize := fs.MustFileSize(chunkPath)
+
+		f, err := os.OpenFile(chunkPath, os.O_WRONLY|os.O_APPEND, 0)
+		if err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+		if _, err := f.Write([]byte("corrupted block")); err != nil {
+			t.Fatalf("cannot append corrupted tail: %s", err)
+		}
+		if err := f.Close(); err != nil {
+			t.Fatalf("cannot close chunk file: %s", err)
+		}
+
+		q = mustOpen(path, "foobar", 0)
+
+		var buf []byte
+		var ok bool
+		for _, block := range blocks {
+			buf, ok = q.MustReadBlockNonblocking(buf[:0])
+			if !ok {
+				t.Fatalf("expected block to be readable")
+			}
+			if string(buf) != string(block) {
+				t.Fatalf("unexpected block got: %q, want: %q", buf, block)
+			}
+		}
+		q.MustClose()
+		chunkFileSize := fs.MustFileSize(chunkPath)
+		if chunkFileSize != originChunkFileSize {
+			t.Fatalf("unexpected chunk file size: got %d; want %d", chunkFileSize, originChunkFileSize)
+		}
+		fs.MustRemoveDir(path)
+	})
+	t.Run("damaged-writer-file-partial-header", func(t *testing.T) {
+		path := "damaged-writer-file-partial-header"
+		fs.MustRemoveDir(path)
+
+		q := mustOpen(path, "foobar", 0)
+		block1 := []byte("valid block 1")
+		block2 := []byte("valid block 2")
+		q.MustWriteBlock(block1)
+		q.MustWriteBlock(block2)
+		q.MustClose()
+
+		chunkPath := filepath.Join(path, fmt.Sprintf("%016X", 0))
+		// Keep block1 intact plus only 3 bytes of the second block's header.
+		newSize := int64(blockHeaderSize + len(block1) + 3)
+		if err := os.Truncate(chunkPath, newSize); err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+
+		q = mustOpen(path, "foobar", 0)
+		buf, ok := q.MustReadBlockNonblocking(nil)
+		if !ok {
+			t.Fatalf("expected block to be readable")
+		}
+		if string(buf) != string(block1) {
+			t.Fatalf("unexpected block got: %q, want: %q", buf, block1)
+		}
+		if _, ok := q.MustReadBlockNonblocking(buf[:0]); ok {
+			t.Fatalf("expected second block to be dropped")
+		}
+		q.MustClose()
+
+		expectedChunkFileSize := uint64(blockHeaderSize + len(block1))
+		if chunkFileSize := fs.MustFileSize(chunkPath); chunkFileSize != expectedChunkFileSize {
+			t.Fatalf("unexpected chunk file size: got %d; want %d", chunkFileSize, expectedChunkFileSize)
+		}
+		fs.MustRemoveDir(path)
+	})
+	t.Run("damaged-writer-file-corrupted-header", func(t *testing.T) {
+		path := "damaged-writer-file-corrupted-header"
+		fs.MustRemoveDir(path)
+
+		q := mustOpen(path, "foobar", 0)
+		block1 := []byte("valid block 1")
+		block2 := []byte("valid block 2")
+		block3 := []byte("valid block 3")
+		q.MustWriteBlock(block1)
+		q.MustWriteBlock(block2)
+		q.MustWriteBlock(block3)
+		q.MustClose()
+
+		chunkPath := filepath.Join(path, fmt.Sprintf("%016X", 0))
+		chunkFileSize := fs.MustFileSize(chunkPath)
+
+		// Overwrite the second block's header with an impossible blockLen
+		// and drop the last byte so the recovery scan is triggered.
+		f, err := os.OpenFile(chunkPath, os.O_WRONLY, 0)
+		if err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+		garbage := [blockHeaderSize]byte{0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF}
+		if _, err := f.WriteAt(garbage[:], int64(blockHeaderSize+len(block1))); err != nil {
+			t.Fatalf("cannot corrupt block header: %s", err)
+		}
+		if err := f.Close(); err != nil {
+			t.Fatalf("cannot close chunk file: %s", err)
+		}
+		if err := os.Truncate(chunkPath, int64(chunkFileSize)-1); err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+
+		q = mustOpen(path, "foobar", 0)
+		buf, ok := q.MustReadBlockNonblocking(nil)
+		if !ok {
+			t.Fatalf("expected block to be readable")
+		}
+		if string(buf) != string(block1) {
+			t.Fatalf("unexpected block got: %q, want: %q", buf, block1)
+		}
+		if _, ok := q.MustReadBlockNonblocking(buf[:0]); ok {
+			t.Fatalf("expected remaining blocks to be dropped")
+		}
+		q.MustClose()
+
+		expectedChunkFileSize := uint64(blockHeaderSize + len(block1))
+		if chunkFileSize := fs.MustFileSize(chunkPath); chunkFileSize != expectedChunkFileSize {
+			t.Fatalf("unexpected chunk file size: got %d; want %d", chunkFileSize, expectedChunkFileSize)
+		}
+		fs.MustRemoveDir(path)
+	})
+	t.Run("damaged-writer-file-second-chunk", func(t *testing.T) {
+		path := "damaged-writer-file-second-chunk"
+		fs.MustRemoveDir(path)
+
+		const maxBlockSize = 64
+		const chunkFileSize = (maxBlockSize + blockHeaderSize) * 2
+
+		newBlock := func(fill byte) []byte {
+			b := make([]byte, 32)
+			for i := range b {
+				b[i] = fill
+			}
+			return b
+		}
+		block1 := newBlock('1')
+		block2 := newBlock('2')
+		block3 := newBlock('3')
+		block4 := newBlock('4')
+
+		q := mustOpenInternal(path, "foobar", chunkFileSize, maxBlockSize, 0)
+		q.MustWriteBlock(block1)
+		q.MustWriteBlock(block2)
+		// blocks 3 and 4 land in the second chunk file
+		q.MustWriteBlock(block3)
+		q.MustWriteBlock(block4)
+		q.MustClose()
+
+		secondChunkPath := filepath.Join(path, fmt.Sprintf("%016X", chunkFileSize))
+		secondChunkSize := fs.MustFileSize(secondChunkPath)
+		if err := os.Truncate(secondChunkPath, int64(secondChunkSize)-1); err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+
+		q = mustOpenInternal(path, "foobar", chunkFileSize, maxBlockSize, 0)
+		var buf []byte
+		var ok bool
+		for _, block := range [][]byte{block1, block2, block3} {
+			buf, ok = q.MustReadBlockNonblocking(buf[:0])
+			if !ok {
+				t.Fatalf("expected block to be readable")
+			}
+			if string(buf) != string(block) {
+				t.Fatalf("unexpected block got: %q, want: %q", buf, block)
+			}
+		}
+		if _, ok := q.MustReadBlockNonblocking(buf[:0]); ok {
+			t.Fatalf("expected last block to be dropped")
+		}
+		q.MustClose()
+
+		expectedSize := uint64(blockHeaderSize + len(block3))
+		if gotSize := fs.MustFileSize(secondChunkPath); gotSize != expectedSize {
+			t.Fatalf("unexpected second chunk file size: got %d; want %d", gotSize, expectedSize)
+		}
+		fs.MustRemoveDir(path)
+	})
 }
 
 func TestQueueResetIfEmpty(t *testing.T) {