From be310e84de07608852156cbb388f1df9712a436d Mon Sep 17 00:00:00 2001
From: Roman Khafizianov <requilence@gmail.com>
Date: Tue, 25 Jun 2024 15:37:08 +0200
Subject: [PATCH] handle badger iterator panics on app stop

---
 .../localstore/objectstore/badger_helpers.go  | 20 ++++++++++++++++---
 pkg/lib/localstore/objectstore/queries.go     |  9 ++++++++-
 pkg/lib/localstore/stores.go                  |  5 +++++
 util/persistentqueue/storage.go               |  8 +++++++-
 4 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/pkg/lib/localstore/objectstore/badger_helpers.go b/pkg/lib/localstore/objectstore/badger_helpers.go
index 4d5c7eb54e..8c16f93689 100644
--- a/pkg/lib/localstore/objectstore/badger_helpers.go
+++ b/pkg/lib/localstore/objectstore/badger_helpers.go
@@ -1,6 +1,8 @@
 package objectstore
 
 import (
+	"fmt"
+
 	"github.com/dgraph-io/badger/v4"
 
 	"github.com/anyproto/anytype-heart/util/badgerhelper"
@@ -35,6 +37,12 @@ func iterateKeysByPrefixBatchedTx(
 	batchSize int,
 	processKeysFn func(keys [][]byte) error,
 ) error {
+	var err error
+	defer func() {
+		if r := recover(); r != nil {
+			err = fmt.Errorf("badger iterator panic: %v", r)
+		}
+	}()
 	opts := badger.DefaultIteratorOptions
 	opts.PrefetchValues = false
 	opts.Prefix = prefix
@@ -50,7 +58,7 @@ func iterateKeysByPrefixBatchedTx(
 		count++
 
 		if count == batchSize {
-			err := processKeysFn(batch)
+			err = processKeysFn(batch)
 			if err != nil {
 				return err
 			}
@@ -60,7 +68,7 @@ func iterateKeysByPrefixBatchedTx(
 	}
 
 	if count > 0 {
-		err := processKeysFn(batch)
+		err = processKeysFn(batch)
 		if err != nil {
 			return err
 		}
@@ -70,6 +78,12 @@ func iterateKeysByPrefixBatchedTx(
 }
 
 func iterateKeysByPrefixTx(txn *badger.Txn, prefix []byte, processKeyFn func(key []byte)) error {
+	var err error
+	defer func() {
+		if r := recover(); r != nil {
+			err = fmt.Errorf("badger iterator panic: %v", r)
+		}
+	}()
 	opts := badger.DefaultIteratorOptions
 	opts.PrefetchValues = false
 	opts.Prefix = prefix
@@ -80,5 +94,5 @@ func iterateKeysByPrefixTx(txn *badger.Txn, prefix []byte, processKeyFn func(key
 		key := iter.Item().KeyCopy(nil)
 		processKeyFn(key)
 	}
-	return nil
+	return err
 }
diff --git a/pkg/lib/localstore/objectstore/queries.go b/pkg/lib/localstore/objectstore/queries.go
index 1ad0980da9..f00042959b 100644
--- a/pkg/lib/localstore/objectstore/queries.go
+++ b/pkg/lib/localstore/objectstore/queries.go
@@ -93,9 +93,16 @@ func (s *dsObjectStore) getInjectedResults(details *types.Struct, score float64,
 func (s *dsObjectStore) queryRaw(filter func(g *types.Struct) bool, order database.Order, limit int, offset int) ([]database.Record, error) {
 	var (
 		records []database.Record
+		err     error
 	)
 
-	err := s.db.View(func(txn *badger.Txn) error {
+	defer func() {
+		if r := recover(); r != nil {
+			err = fmt.Errorf("badger iterator panic: %v", r)
+		}
+	}()
+
+	err = s.db.View(func(txn *badger.Txn) error {
 		opts := badger.DefaultIteratorOptions
 		opts.PrefetchValues = false
 		opts.Prefix = pagesDetailsBase.Bytes()
diff --git a/pkg/lib/localstore/stores.go b/pkg/lib/localstore/stores.go
index 124b99d62d..79fb55ee19 100644
--- a/pkg/lib/localstore/stores.go
+++ b/pkg/lib/localstore/stores.go
@@ -238,6 +238,11 @@ func GetKeysByIndex(index Index, txn *badger.Txn, val interface{}, limit int) ([
 }
 
 func GetKeys(txn *badger.Txn, prefix string, limit int) []string {
+	defer func() {
+		if r := recover(); r != nil {
+			log.Errorf("badger iterator panic: %v", r)
+		}
+	}()
 	iter := txn.NewIterator(badger.IteratorOptions{
 		Prefix:         []byte(prefix),
 		PrefetchValues: false,
diff --git a/util/persistentqueue/storage.go b/util/persistentqueue/storage.go
index 52c39787af..1b4910925f 100644
--- a/util/persistentqueue/storage.go
+++ b/util/persistentqueue/storage.go
@@ -34,7 +34,13 @@ func NewBadgerStorage[T Item](db *badger.DB, badgerPrefix []byte, factoryFunc Fa
 
 func (s *badgerStorage[T]) List() ([]T, error) {
 	var items []T
-	err := s.db.View(func(txn *badger.Txn) error {
+	var err error
+	defer func() {
+		if r := recover(); r != nil {
+			err = fmt.Errorf("badger iterator panic: %v", r)
+		}
+	}()
+	err = s.db.View(func(txn *badger.Txn) error {
 		it := txn.NewIterator(badger.IteratorOptions{
 			PrefetchSize:   100,
 			PrefetchValues: true,