From cfee8d550aa0779721d85ee2e2935fed651d79df Mon Sep 17 00:00:00 2001
From: Sendya <18x@loacg.com>
Date: Thu, 19 Mar 2026 14:58:28 +0800
Subject: [PATCH] fix: PebbleDB and noneSharedKV Close idempotent using
 `sync.Once`. fix hot upgrade process

---
 storage/indexdb/pebble/pebble.go | 13 ++++++++++---
 storage/sharedkv/nonekv.go       | 13 ++++++++++---
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/storage/indexdb/pebble/pebble.go b/storage/indexdb/pebble/pebble.go
index 7974506..7ebb4a6 100644
--- a/storage/indexdb/pebble/pebble.go
+++ b/storage/indexdb/pebble/pebble.go
@@ -3,6 +3,7 @@ package pebble
 import (
 	"context"
 	"errors"
+	"sync"
 	"time"
 
 	"github.com/cockroachdb/pebble/v2"
@@ -22,6 +23,7 @@ type PebbleDB struct {
 	db            *pebble.DB
 	writeMode     *pebble.WriteOptions
 	skipErrRecord bool
+	closed        sync.Once
 }
 
 func init() {
@@ -112,9 +114,13 @@ func (p *PebbleDB) Expired(ctx context.Context, f storage.IterateFunc) error {
 
 // Close implements storage.IndexDB.
 func (p *PebbleDB) Close() error {
-	// force flush data to disk
-	_ = p.db.Flush()
-	return p.db.Close()
+	var err error
+	p.closed.Do(func() {
+		_ = p.db.Flush()
+		// force flush data to disk
+		err = p.db.Close()
+	})
+	return err
 }
 
 // GC implements storage.IndexDB.
@@ -178,5 +184,6 @@ func New(path string, option storage.Option) (storage.IndexDB, error) {
 		db:            pdb,
 		writeMode:     writeMode, // 是否异步写操作
 		skipErrRecord: true,
+		closed:        sync.Once{},
 	}, nil
 }
diff --git a/storage/sharedkv/nonekv.go b/storage/sharedkv/nonekv.go
index 8dcec72..0eaf90c 100644
--- a/storage/sharedkv/nonekv.go
+++ b/storage/sharedkv/nonekv.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/binary"
 	"errors"
+	"sync"
 
 	"github.com/cockroachdb/pebble/v2"
 	"github.com/omalloc/tavern/api/defined/v1/storage"
@@ -12,11 +13,16 @@ import (
 var _ storage.SharedKV = (*noneSharedKV)(nil)
 
 type noneSharedKV struct {
-	db *pebble.DB
+	db     *pebble.DB
+	closed sync.Once
 }
 
 func (r *noneSharedKV) Close() error {
-	return r.db.Close()
+	var err error
+	r.closed.Do(func() {
+		err = r.db.Close()
+	})
+	return err
 }
 
 func (r *noneSharedKV) Get(_ context.Context, key []byte) ([]byte, error) {
@@ -189,7 +195,8 @@ func newNoneKV(storePath string, opts *pebble.Options) (storage.SharedKV, error)
 	}
 
 	r := &noneSharedKV{
-		db: db,
+		db:     db,
+		closed: sync.Once{},
 	}
 	return r, nil
 }