From 7a8b33cb0c52221836a0c69f8e141705bbdaf627 Mon Sep 17 00:00:00 2001
From: Anton Kolesnikov <anton.e.kolesnikov@gmail.com>
Date: Tue, 18 Oct 2022 16:27:52 +0200
Subject: [PATCH] fix: close databases in deterministic order (#1623)

---
 pkg/storage/db.go      |  9 ------
 pkg/storage/storage.go | 64 ++++++++++++++++++++++++++++--------------
 pkg/storage/types.go   |  1 -
 3 files changed, 43 insertions(+), 31 deletions(-)

diff --git a/pkg/storage/db.go b/pkg/storage/db.go
index 47699632b9..5ac03bea7c 100644
--- a/pkg/storage/db.go
+++ b/pkg/storage/db.go
@@ -134,15 +134,6 @@ func (s *Storage) newBadger(name string, p Prefix, codec cache.Codec) (BadgerDBW
 	return d, nil
 }
 
-func (d *db) Close() {
-	if d.Cache != nil {
-		d.Cache.Flush()
-	}
-	if err := d.DB.Close(); err != nil {
-		d.logger.WithError(err).Error("closing database")
-	}
-}
-
 func (d *db) Size() bytesize.ByteSize {
 	// The value is updated once per minute.
 	lsm, vlog := d.DB.Size()
diff --git a/pkg/storage/storage.go b/pkg/storage/storage.go
index 8552d7755c..2baaa099cb 100644
--- a/pkg/storage/storage.go
+++ b/pkg/storage/storage.go
@@ -161,13 +161,49 @@ func (s *Storage) Close() error {
 	s.logger.Debug("waiting for storage tasks to finish")
 	s.tasksWG.Wait()
 	s.logger.Debug("storage tasks finished")
-	// Dictionaries DB has to close last because trees and profiles DBs depend on it.
-	s.goDB(func(d BadgerDBWithCache) {
-		if d != s.dicts {
-			d.Close()
-		}
-	})
-	s.dicts.Close()
+
+	// Flush caches. Dictionaries DB has to close last because trees depends on it.
+	// Exemplars DB does not have a cache but depends on Dictionaries DB as well:
+	// there is no need to force synchronization, as exemplars storage listens to
+	// the s.stop channel and stops synchronously.
+	caches := []BadgerDBWithCache{
+		s.trees,
+		s.segments,
+		s.dimensions,
+	}
+	wg := new(sync.WaitGroup)
+	wg.Add(len(caches))
+	for _, d := range caches {
+		go func(d BadgerDBWithCache) {
+			d.CacheInstance().Flush()
+			wg.Done()
+		}(d)
+	}
+	wg.Wait()
+
+	// Flush dictionaries cache only when all the dependant caches are flushed.
+	s.dicts.CacheInstance().Flush()
+
+	// Close databases. Order does not matter.
+	dbs := []BadgerDBWithCache{
+		s.trees,
+		s.segments,
+		s.dimensions,
+		s.exemplars.db,
+		s.dicts,
+		s.main, // Also stores labels.
+	}
+	wg = new(sync.WaitGroup)
+	wg.Add(len(dbs))
+	for _, d := range dbs {
+		go func(d BadgerDBWithCache) {
+			defer wg.Done()
+			if err := d.DBInstance().Close(); err != nil {
+				s.logger.WithField("name", d.Name()).WithError(err).Error("closing database")
+			}
+		}(d)
+	}
+	wg.Wait()
 	return nil
 }
 
@@ -202,20 +238,6 @@ func (s *Storage) withContext(fn func(context.Context)) {
 	fn(ctx)
 }
 
-// goDB runs f for all DBs concurrently.
-func (s *Storage) goDB(f func(BadgerDBWithCache)) {
-	dbs := s.databases()
-	wg := new(sync.WaitGroup)
-	wg.Add(len(dbs))
-	for _, d := range dbs {
-		go func(db BadgerDBWithCache) {
-			defer wg.Done()
-			f(db)
-		}(d)
-	}
-	wg.Wait()
-}
-
 // maintenanceTask periodically runs f exclusively.
 func (s *Storage) maintenanceTask(interval time.Duration, f func()) {
 	s.periodicTask(interval, func() {
diff --git a/pkg/storage/types.go b/pkg/storage/types.go
index 282de054fc..ab25885f96 100644
--- a/pkg/storage/types.go
+++ b/pkg/storage/types.go
@@ -124,7 +124,6 @@ type BadgerDBWithCache interface {
 	BadgerDB
 	CacheLayer
 
-	Close()
 	Size() bytesize.ByteSize
 	CacheSize() uint64