From 30534e99d97630d3d6acf804f536c474076c9f67 Mon Sep 17 00:00:00 2001
From: Ganesh Vernekar <15064823+codesome@users.noreply.github.com>
Date: Mon, 13 Sep 2021 18:30:41 +0530
Subject: [PATCH] Take snapshot only after closing the WAL (#9328)

Signed-off-by: Ganesh Vernekar <ganeshvern@gmail.com>
---
 tsdb/head.go      | 6 +++---
 tsdb/head_test.go | 8 ++++++++
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/tsdb/head.go b/tsdb/head.go
index 8d35f99b8..ec0916bba 100644
--- a/tsdb/head.go
+++ b/tsdb/head.go
@@ -1168,12 +1168,12 @@ func (h *Head) Close() error {
 	defer h.closedMtx.Unlock()
 	h.closed = true
 	errs := tsdb_errors.NewMulti(h.chunkDiskMapper.Close())
-	if errs.Err() == nil && h.opts.EnableMemorySnapshotOnShutdown {
-		errs.Add(h.performChunkSnapshot())
-	}
 	if h.wal != nil {
 		errs.Add(h.wal.Close())
 	}
+	if errs.Err() == nil && h.opts.EnableMemorySnapshotOnShutdown {
+		errs.Add(h.performChunkSnapshot())
+	}
 	return errs.Err()
 
 }
diff --git a/tsdb/head_test.go b/tsdb/head_test.go
index b09e59ae9..7d4b63843 100644
--- a/tsdb/head_test.go
+++ b/tsdb/head_test.go
@@ -2659,6 +2659,10 @@ func TestChunkSnapshot(t *testing.T) {
 	// These references should be the ones used for the snapshot.
 	wlast, woffset, err = head.wal.LastSegmentAndOffset()
 	require.NoError(t, err)
+	if woffset != 0 && woffset < 32*1024 {
+		// The page is always filled before taking the snapshot.
+		woffset = 32 * 1024
+	}
 
 	{
 		// Creating snapshot and verifying it.
@@ -2725,6 +2729,10 @@ func TestChunkSnapshot(t *testing.T) {
 	// Creating another snapshot should delete the older snapshot and replay still works fine.
 	wlast, woffset, err = head.wal.LastSegmentAndOffset()
 	require.NoError(t, err)
+	if woffset != 0 && woffset < 32*1024 {
+		// The page is always filled before taking the snapshot.
+		woffset = 32 * 1024
+	}
 
 	{
 		// Close Head and verify that new snapshot was created.