From dbea2c275e3dca167971a4b8bc79f82774ffd5b2 Mon Sep 17 00:00:00 2001
From: Dmitrii Stepanov <d.stepanov@yadro.com>
Date: Thu, 29 Aug 2024 16:33:30 +0300
Subject: [PATCH] [#1334] metabase: Store upgrade flag

This allows to check if metabase upgrade was not completed.

Signed-off-by: Dmitrii Stepanov <d.stepanov@yadro.com>
---
 pkg/local_object_storage/metabase/upgrade.go  | 12 +++++++++++
 pkg/local_object_storage/metabase/version.go  | 11 +++++++++-
 .../metabase/version_test.go                  | 20 +++++++++++++++++++
 pkg/local_object_storage/shard/control.go     |  2 +-
 4 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/pkg/local_object_storage/metabase/upgrade.go b/pkg/local_object_storage/metabase/upgrade.go
index 014e50286..a4c7707b4 100644
--- a/pkg/local_object_storage/metabase/upgrade.go
+++ b/pkg/local_object_storage/metabase/upgrade.go
@@ -51,9 +51,21 @@ func Upgrade(ctx context.Context, path string, compact bool, log func(a ...any))
 	if !found {
 		return fmt.Errorf("unsupported version %d: no update available", version)
 	}
+	if err := db.Update(func(tx *bbolt.Tx) error {
+		b := tx.Bucket(shardInfoBucket)
+		return b.Put(upgradeKey, zeroValue)
+	}); err != nil {
+		return fmt.Errorf("set upgrade key %w", err)
+	}
 	if err := updater(ctx, db, log); err != nil {
 		return fmt.Errorf("update metabase schema: %w", err)
 	}
+	if err := db.Update(func(tx *bbolt.Tx) error {
+		b := tx.Bucket(shardInfoBucket)
+		return b.Delete(upgradeKey)
+	}); err != nil {
+		return fmt.Errorf("delete upgrade key %w", err)
+	}
 	if compact {
 		log("compacting metabase...")
 		err := compactDB(db)
diff --git a/pkg/local_object_storage/metabase/version.go b/pkg/local_object_storage/metabase/version.go
index 9e15babbc..048bb9af6 100644
--- a/pkg/local_object_storage/metabase/version.go
+++ b/pkg/local_object_storage/metabase/version.go
@@ -12,13 +12,18 @@ import (
 // version contains current metabase version.
 const version = 3
 
-var versionKey = []byte("version")
+var (
+	versionKey = []byte("version")
+	upgradeKey = []byte("upgrade")
+)
 
 // ErrOutdatedVersion is returned on initializing
 // an existing metabase that is not compatible with
 // the current code version.
 var ErrOutdatedVersion = logicerr.New("invalid version, resynchronization is required")
 
+var ErrIncompletedUpgrade = logicerr.New("metabase upgrade is not completed")
+
 var errVersionUndefinedNoInfoBucket = errors.New("version undefined: no info bucket")
 
 func checkVersion(tx *bbolt.Tx, initialized bool) error {
@@ -35,6 +40,10 @@ func checkVersion(tx *bbolt.Tx, initialized bool) error {
 				return fmt.Errorf("%w: expected=%d, stored=%d", ErrOutdatedVersion, version, stored)
 			}
 		}
+		data = b.Get(upgradeKey)
+		if len(data) > 0 {
+			return ErrIncompletedUpgrade
+		}
 	}
 
 	if !initialized {
diff --git a/pkg/local_object_storage/metabase/version_test.go b/pkg/local_object_storage/metabase/version_test.go
index b2af428ff..75229a1b4 100644
--- a/pkg/local_object_storage/metabase/version_test.go
+++ b/pkg/local_object_storage/metabase/version_test.go
@@ -84,4 +84,24 @@ func TestVersion(t *testing.T) {
 			require.NoError(t, db.Close())
 		})
 	})
+	t.Run("incompleted upgrade", func(t *testing.T) {
+		db := newDB(t)
+		require.NoError(t, db.Open(context.Background(), mode.ReadWrite))
+		require.NoError(t, db.Init())
+		require.NoError(t, db.Close())
+
+		require.NoError(t, db.Open(context.Background(), mode.ReadWrite))
+		require.NoError(t, db.boltDB.Update(func(tx *bbolt.Tx) error {
+			return tx.Bucket(shardInfoBucket).Put(upgradeKey, zeroValue)
+		}))
+		require.ErrorIs(t, db.Init(), ErrIncompletedUpgrade)
+		require.NoError(t, db.Close())
+
+		require.NoError(t, db.Open(context.Background(), mode.ReadWrite))
+		require.NoError(t, db.boltDB.Update(func(tx *bbolt.Tx) error {
+			return tx.Bucket(shardInfoBucket).Delete(upgradeKey)
+		}))
+		require.NoError(t, db.Init())
+		require.NoError(t, db.Close())
+	})
 }
diff --git a/pkg/local_object_storage/shard/control.go b/pkg/local_object_storage/shard/control.go
index 936a506c0..1626d5804 100644
--- a/pkg/local_object_storage/shard/control.go
+++ b/pkg/local_object_storage/shard/control.go
@@ -171,7 +171,7 @@ func (s *Shard) initializeComponents(m mode.Mode) error {
 	for _, component := range components {
 		if err := component.Init(); err != nil {
 			if component == s.metaBase {
-				if errors.Is(err, meta.ErrOutdatedVersion) {
+				if errors.Is(err, meta.ErrOutdatedVersion) || errors.Is(err, meta.ErrIncompletedUpgrade) {
 					return fmt.Errorf("metabase initialization: %w", err)
 				}