refactor(repository): refactored Repository close logic (#2589)

This ensures metric registry is always closed on last Close().

internal/metrics/metrics_registry.go

@@ -7,6 +7,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/kopia/kopia/internal/releasable"
 	"github.com/kopia/kopia/repo/logging"
 )
 
@@ -64,6 +65,8 @@ func (r *Registry) Close(ctx context.Context) error {
 		return nil
 	}
 
+	releasable.Released("metric-registry", r)
+
 	return nil
 }
 
@@ -98,13 +101,17 @@ func (r *Registry) Log(ctx context.Context) error {
 
 // NewRegistry returns new metrics registry.
 func NewRegistry() *Registry {
-	return &Registry{
+	r := &Registry{
 		allCounters:              map[string]*Counter{},
 		allGauges:                map[string]*Gauge{},
 		allDurationDistributions: map[string]*Distribution[time.Duration]{},
 		allSizeDistributions:     map[string]*Distribution[int64]{},
 		allThroughput:            map[string]*Throughput{},
 	}
+
+	releasable.Created("metric-registry", r)
+
+	return r
 }
 
 func labelsSuffix(l map[string]string) string {

repo/api_server_repository.go

@@ -44,6 +44,7 @@ type apiServerRepository struct {
 
 	isSharedReadOnlySession bool
 	contentCache            *cache.PersistentCache
+	*refCountedCloser
 }
 
 func (r *apiServerRepository) APIServerURL() string {
@@ -164,7 +165,8 @@ func (r *apiServerRepository) NewWriter(ctx context.Context, opt WriteSessionOpt
 
 	w.omgr = omgr
 	w.wso = opt
-	w.isSharedReadOnlySession = false
+
+	r.refCountedCloser.addRef()
 
 	return ctx, w, nil
 }
@@ -247,15 +249,6 @@ func (r *apiServerRepository) UpdateDescription(d string) {
 	r.cliOpts.Description = d
 }
 
-func (r *apiServerRepository) Close(ctx context.Context) error {
-	if r.isSharedReadOnlySession && r.contentCache != nil {
-		r.contentCache.Close(ctx)
-		r.contentCache = nil
-	}
-
-	return nil
-}
-
 func (r *apiServerRepository) PrefetchObjects(ctx context.Context, objectIDs []object.ID, hint string) ([]content.ID, error) {
 	//nolint:wrapcheck
 	return object.PrefetchBackingContents(ctx, r, objectIDs, hint)
@@ -295,6 +288,17 @@ func openRestAPIRepository(ctx context.Context, si *APIServerInfo, cliOpts Clien
 		return nil, errors.Wrap(err, "unable to create API client")
 	}
 
+	rcc := newRefCountedCloser(
+		func(ctx context.Context) error {
+			if contentCache != nil {
+				contentCache.Close(ctx)
+			}
+
+			return nil
+		},
+		mr.Close,
+	)
+
 	rr := &apiServerRepository{
 		cli:          cli,
 		cliOpts:      cliOpts,
@@ -304,6 +308,7 @@ func openRestAPIRepository(ctx context.Context, si *APIServerInfo, cliOpts Clien
 		},
 		isSharedReadOnlySession: true,
 		metricsRegistry:         mr,
+		refCountedCloser:        rcc,
 	}
 
 	var p remoterepoapi.Parameters

repo/content/committed_read_manager.go

@@ -5,7 +5,6 @@ import (
 	"os"
 	"path/filepath"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/pkg/errors"
@@ -74,11 +73,6 @@ type indexBlobManager interface {
 
 // SharedManager is responsible for read-only access to committed data.
 type SharedManager struct {
-	// +checkatomic
-	refCount int32 // number of Manager objects that refer to this SharedManager
-	// +checkatomic
-	closed int32 // set to 1 if shared manager has been closed
-
 	Stats *Stats
 	st    blob.Storage
 
@@ -520,33 +514,8 @@ func (sm *SharedManager) EpochManager() (*epoch.Manager, bool, error) {
 	return ibm1.epochMgr, true, nil
 }
 
-// AddRef adds a reference to shared manager to prevents its closing on Release().
-func (sm *SharedManager) addRef() {
-	if atomic.LoadInt32(&sm.closed) != 0 {
-		panic("attempted to re-use closed SharedManager")
-	}
-
-	atomic.AddInt32(&sm.refCount, 1)
-}
-
-// release removes a reference to the shared manager and destroys it if no more references are remaining.
-func (sm *SharedManager) release(ctx context.Context) error {
-	if atomic.LoadInt32(&sm.closed) != 0 {
-		// already closed
-		return nil
-	}
-
-	remaining := atomic.AddInt32(&sm.refCount, -1)
-	if remaining != 0 {
-		sm.log.Debugf("not closing shared manager, remaining = %v", remaining)
-
-		return nil
-	}
-
-	atomic.StoreInt32(&sm.closed, 1)
-
-	sm.log.Debugf("closing shared manager")
-
+// CloseShared releases all resources in a shared manager.
+func (sm *SharedManager) CloseShared(ctx context.Context) error {
 	if err := sm.committedContents.close(); err != nil {
 		return errors.Wrap(err, "error closing committed content index")
 	}
@@ -567,10 +536,6 @@ func (sm *SharedManager) release(ctx context.Context) error {
 		return errors.Wrap(err, "error closing storage")
 	}
 
-	if err := sm.metricsEmitter.Close(ctx); err != nil {
-		return errors.Wrap(err, "error closing metrics")
-	}
-
 	return nil
 }
 

repo/content/content_formatter_test.go

@@ -93,7 +93,7 @@ func verifyEndToEndFormatter(ctx context.Context, t *testing.T, hashAlgo, encryp
 		return
 	}
 
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 
 	cases := []gather.Bytes{
 		gather.FromSlice([]byte{}),

repo/content/content_index_recovery_test.go

@@ -38,11 +38,11 @@ func (s *contentManagerSuite) TestContentIndexRecovery(t *testing.T) {
 		return bm.st.DeleteBlob(ctx, bi.BlobID)
 	}))
 
-	bm.Close(ctx)
+	bm.CloseShared(ctx)
 
 	// now with index blobs gone, all contents appear to not be found
 	bm = s.newTestContentManagerWithCustomTime(t, st, nil)
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 
 	verifyContentNotFound(ctx, t, bm, content1)
 	verifyContentNotFound(ctx, t, bm, content2)

repo/content/content_manager.go

@@ -589,11 +589,6 @@ func (bm *WriteManager) ContentFormat() format.Provider {
 	return bm.format
 }
 
-// Close closes the content manager.
-func (bm *WriteManager) Close(ctx context.Context) error {
-	return bm.SharedManager.release(ctx)
-}
-
 // +checklocks:bm.mu
 func (bm *WriteManager) setFlushingLocked(v bool) {
 	bm.flushing = v
@@ -996,8 +991,6 @@ type SessionOptions struct {
 
 // NewWriteManager returns a session write manager.
 func NewWriteManager(ctx context.Context, sm *SharedManager, options SessionOptions, writeManagerID string) *WriteManager {
-	sm.addRef()
-
 	if options.OnUpload == nil {
 		options.OnUpload = func(int64) {}
 	}

repo/content/content_manager_metrics.go

@@ -5,8 +5,6 @@ import (
 )
 
 type metricsStruct struct {
-	metricsEmitter *metrics.Registry
-
 	uploadedBytes *metrics.Counter
 
 	getContentBytes         *metrics.Throughput
@@ -36,8 +34,6 @@ type metricsStruct struct {
 
 func initMetricsStruct(mr *metrics.Registry) metricsStruct {
 	return metricsStruct{
-		metricsEmitter: mr,
-
 		uploadedBytes:           mr.CounterInt64("content_uploaded_bytes", "Number of bytes uploaded from content manager.", nil),
 		getContentErrorCount:    mr.CounterInt64("content_get_error_count", "Number of time GetContent() was called and the result was an error", nil),
 		getContentNotFoundCount: mr.CounterInt64("content_get_not_found_count", "Number of time GetContent() was called and the result was not found", nil),

repo/content/content_manager_test.go

@@ -82,7 +82,7 @@ func (s *contentManagerSuite) TestContentManagerEmptyFlush(t *testing.T) {
 	st := blobtesting.NewMapStorage(data, nil, nil)
 	bm := s.newTestContentManager(t, st)
 
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 	bm.Flush(ctx)
 
 	if got, want := len(data), 0; got != want {
@@ -96,7 +96,7 @@ func (s *contentManagerSuite) TestContentZeroBytes1(t *testing.T) {
 	st := blobtesting.NewMapStorage(data, nil, nil)
 	bm := s.newTestContentManager(t, st)
 
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 	contentID := writeContentAndVerify(ctx, t, bm, []byte{})
 	bm.Flush(ctx)
 
@@ -107,7 +107,7 @@ func (s *contentManagerSuite) TestContentZeroBytes1(t *testing.T) {
 	dumpContentManagerData(t, data)
 	bm = s.newTestContentManager(t, st)
 
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 
 	verifyContent(ctx, t, bm, contentID, []byte{})
 }
@@ -118,7 +118,7 @@ func (s *contentManagerSuite) TestContentZeroBytes2(t *testing.T) {
 	st := blobtesting.NewMapStorage(data, nil, nil)
 	bm := s.newTestContentManager(t, st)
 
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 
 	writeContentAndVerify(ctx, t, bm, seededRandomData(10, 10))
 	writeContentAndVerify(ctx, t, bm, []byte{})
@@ -136,7 +136,7 @@ func (s *contentManagerSuite) TestContentManagerSmallContentWrites(t *testing.T)
 	st := blobtesting.NewMapStorage(data, nil, nil)
 	bm := s.newTestContentManager(t, st)
 
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 
 	itemCount := maxPackCapacity / (10 + encryptionOverhead)
 	for i := 0; i < itemCount; i++ {
@@ -159,7 +159,7 @@ func (s *contentManagerSuite) TestContentManagerDedupesPendingContents(t *testin
 	st := blobtesting.NewMapStorage(data, nil, nil)
 	bm := s.newTestContentManager(t, st)
 
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 
 	for i := 0; i < 100; i++ {
 		writeContentAndVerify(ctx, t, bm, seededRandomData(0, maxPackCapacity/2))
@@ -184,7 +184,7 @@ func (s *contentManagerSuite) TestContentManagerDedupesPendingAndUncommittedCont
 	st := blobtesting.NewMapStorage(data, nil, nil)
 	bm := s.newTestContentManager(t, st)
 
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 
 	// compute content size so that 3 contents will fit in a pack without overflowing
 	contentSize := maxPackCapacity/3 - encryptionOverhead - 1
@@ -222,7 +222,7 @@ func (s *contentManagerSuite) TestContentManagerEmpty(t *testing.T) {
 	st := blobtesting.NewMapStorage(data, nil, nil)
 	bm := s.newTestContentManager(t, st)
 
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 
 	noSuchContentID := hashValue(t, []byte("foo"))
 
@@ -259,7 +259,7 @@ func (s *contentManagerSuite) TestContentManagerInternalFlush(t *testing.T) {
 	st := blobtesting.NewMapStorage(data, nil, nil)
 	bm := s.newTestContentManager(t, st)
 
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 
 	itemsToOverflow := (maxPackCapacity)/(25+encryptionOverhead) + 2
 	for i := 0; i < itemsToOverflow; i++ {
@@ -299,7 +299,7 @@ func (s *contentManagerSuite) TestContentManagerWriteMultiple(t *testing.T) {
 	st := blobtesting.NewMapStorage(data, keyTime, timeFunc)
 
 	bm := s.newTestContentManagerWithCustomTime(t, st, timeFunc)
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 
 	var contentIDs []ID
 
@@ -330,7 +330,7 @@ func (s *contentManagerSuite) TestContentManagerWriteMultiple(t *testing.T) {
 			}
 
 			bm = s.newTestContentManagerWithCustomTime(t, st, timeFunc)
-			defer bm.Close(ctx) //nolint:gocritic
+			defer bm.CloseShared(ctx) //nolint:gocritic
 		}
 
 		pos := rand.Intn(len(contentIDs))
@@ -366,7 +366,7 @@ func (s *contentManagerSuite) TestContentManagerFailedToWritePack(t *testing.T)
 		t.Fatalf("can't create bm: %v", err)
 	}
 
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 
 	sessionPutErr := errors.New("booboo0")
 	firstPutErr := errors.New("booboo1")
@@ -424,7 +424,7 @@ func (s *contentManagerSuite) TestIndexCompactionDropsContent(t *testing.T) {
 	bm := s.newTestContentManagerWithCustomTime(t, st, timeFunc)
 	content1 := writeContentAndVerify(ctx, t, bm, seededRandomData(10, 100))
 	require.NoError(t, bm.Flush(ctx))
-	require.NoError(t, bm.Close(ctx))
+	require.NoError(t, bm.CloseShared(ctx))
 
 	timeFunc()
 
@@ -432,7 +432,7 @@ func (s *contentManagerSuite) TestIndexCompactionDropsContent(t *testing.T) {
 	bm = s.newTestContentManagerWithCustomTime(t, st, timeFunc)
 	deleteContent(ctx, t, bm, content1)
 	require.NoError(t, bm.Flush(ctx))
-	require.NoError(t, bm.Close(ctx))
+	require.NoError(t, bm.CloseShared(ctx))
 
 	timeFunc()
 
@@ -447,7 +447,7 @@ func (s *contentManagerSuite) TestIndexCompactionDropsContent(t *testing.T) {
 		AllIndexes:        true,
 	}))
 	require.NoError(t, bm.Flush(ctx))
-	require.NoError(t, bm.Close(ctx))
+	require.NoError(t, bm.CloseShared(ctx))
 
 	bm = s.newTestContentManagerWithCustomTime(t, st, timeFunc)
 	verifyContentNotFound(ctx, t, bm, content1)
@@ -460,7 +460,7 @@ func (s *contentManagerSuite) TestContentManagerConcurrency(t *testing.T) {
 	st := blobtesting.NewMapStorage(data, keyTime, nil)
 
 	bm := s.newTestContentManagerWithCustomTime(t, st, nil)
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 
 	preexistingContent := writeContentAndVerify(ctx, t, bm, seededRandomData(10, 100))
 	bm.Flush(ctx)
@@ -468,13 +468,13 @@ func (s *contentManagerSuite) TestContentManagerConcurrency(t *testing.T) {
 	dumpContentManagerData(t, data)
 
 	bm1 := s.newTestContentManager(t, st)
-	defer bm1.Close(ctx)
+	defer bm1.CloseShared(ctx)
 
 	bm2 := s.newTestContentManager(t, st)
-	defer bm2.Close(ctx)
+	defer bm2.CloseShared(ctx)
 
 	bm3 := s.newTestContentManagerWithCustomTime(t, st, faketime.AutoAdvance(fakeTime.Add(1), 1*time.Second))
-	defer bm3.Close(ctx)
+	defer bm3.CloseShared(ctx)
 
 	// all bm* can see pre-existing content
 	verifyContent(ctx, t, bm1, preexistingContent, seededRandomData(10, 100))
@@ -512,7 +512,7 @@ func (s *contentManagerSuite) TestContentManagerConcurrency(t *testing.T) {
 
 	// new content manager at this point can see all data.
 	bm4 := s.newTestContentManager(t, st)
-	defer bm4.Close(ctx)
+	defer bm4.CloseShared(ctx)
 
 	verifyContent(ctx, t, bm4, preexistingContent, seededRandomData(10, 100))
 	verifyContent(ctx, t, bm4, sharedContent, seededRandomData(20, 100))
@@ -532,7 +532,7 @@ func (s *contentManagerSuite) TestContentManagerConcurrency(t *testing.T) {
 
 	// new content manager at this point can see all data.
 	bm5 := s.newTestContentManager(t, st)
-	defer bm5.Close(ctx)
+	defer bm5.CloseShared(ctx)
 
 	verifyContent(ctx, t, bm5, preexistingContent, seededRandomData(10, 100))
 	verifyContent(ctx, t, bm5, sharedContent, seededRandomData(20, 100))
@@ -575,7 +575,7 @@ func (s *contentManagerSuite) TestDeleteContent(t *testing.T) {
 	st := blobtesting.NewMapStorage(data, nil, nil)
 	bm := s.newTestContentManager(t, st)
 
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 
 	c1Bytes := seededRandomData(10, 100)
 	content1 := writeContentAndVerify(ctx, t, bm, c1Bytes)
@@ -609,7 +609,7 @@ func (s *contentManagerSuite) TestDeleteContent(t *testing.T) {
 	t.Logf("flushed")
 
 	bm = s.newTestContentManager(t, st)
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 
 	verifyDeletedContentRead(ctx, t, bm, content1, c1Bytes)
 	verifyContentNotFound(ctx, t, bm, content2)
@@ -1024,7 +1024,7 @@ func (s *contentManagerSuite) TestParallelWrites(t *testing.T) {
 	var workerLock sync.RWMutex
 
 	bm := s.newTestContentManagerWithTweaks(t, fs, nil)
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 
 	numWorkers := 8
 	closeWorkers := make(chan bool)
@@ -1129,7 +1129,7 @@ func (s *contentManagerSuite) TestFlushResumesWriters(t *testing.T) {
 	})
 
 	bm := s.newTestContentManagerWithTweaks(t, fs, nil)
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 	first := writeContentAndVerify(ctx, t, bm, []byte{1, 2, 3})
 
 	var second ID
@@ -1186,7 +1186,7 @@ func (s *contentManagerSuite) TestFlushWaitsForAllPendingWriters(t *testing.T) {
 	fs.AddFault(blobtesting.MethodPutBlob).SleepFor(2 * time.Second)
 
 	bm := s.newTestContentManagerWithTweaks(t, fs, nil)
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 
 	// write one content in another goroutine
 	// 'fs' is configured so that blob write takes several seconds to complete.
@@ -1225,7 +1225,7 @@ func (s *contentManagerSuite) verifyAllDataPresent(ctx context.Context, t *testi
 	t.Helper()
 
 	bm := s.newTestContentManagerWithCustomTime(t, st, nil)
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 	_ = bm.IterateContents(ctx, IterateOptions{}, func(ci Info) error {
 		delete(contentIDs, ci.GetContentID())
 		return nil
@@ -1297,7 +1297,7 @@ func (s *contentManagerSuite) TestHandleWriteErrors(t *testing.T) {
 			fs.AddFaults(blobtesting.MethodPutBlob, tc.faults...)
 
 			bm := s.newTestContentManagerWithTweaks(t, fs, nil)
-			defer bm.Close(ctx)
+			defer bm.CloseShared(ctx)
 
 			var writeRetries []int
 			var cids []ID
@@ -1315,7 +1315,7 @@ func (s *contentManagerSuite) TestHandleWriteErrors(t *testing.T) {
 			}
 
 			bm2 := s.newTestContentManagerWithTweaks(t, st, nil)
-			defer bm2.Close(ctx)
+			defer bm2.CloseShared(ctx)
 
 			for i, cid := range cids {
 				verifyContent(ctx, t, bm2, cid, seededRandomData(i, tc.contentSizes[i]))
@@ -1344,7 +1344,7 @@ func (s *contentManagerSuite) TestRewriteNonDeleted(t *testing.T) {
 				st := blobtesting.NewMapStorage(data, keyTime, fakeNow)
 
 				bm := s.newTestContentManagerWithCustomTime(t, st, fakeNow)
-				defer bm.Close(ctx)
+				defer bm.CloseShared(ctx)
 
 				applyStep := func(action int) {
 					switch action {
@@ -1352,7 +1352,7 @@ func (s *contentManagerSuite) TestRewriteNonDeleted(t *testing.T) {
 						t.Logf("flushing and reopening")
 						bm.Flush(ctx)
 						bm = s.newTestContentManagerWithCustomTime(t, st, fakeNow)
-						defer bm.Close(ctx)
+						defer bm.CloseShared(ctx)
 					case 1:
 						t.Logf("flushing")
 						bm.Flush(ctx)
@@ -1414,7 +1414,7 @@ func (s *contentManagerSuite) TestRewriteDeleted(t *testing.T) {
 					fakeNow := faketime.AutoAdvance(fakeTime, 1*time.Second)
 					st := blobtesting.NewMapStorage(data, keyTime, fakeNow)
 					bm := s.newTestContentManagerWithCustomTime(t, st, fakeNow)
-					defer bm.Close(ctx)
+					defer bm.CloseShared(ctx)
 
 					applyStep := func(action int) {
 						switch action {
@@ -1422,7 +1422,7 @@ func (s *contentManagerSuite) TestRewriteDeleted(t *testing.T) {
 							t.Logf("flushing and reopening")
 							bm.Flush(ctx)
 							bm = s.newTestContentManagerWithCustomTime(t, st, fakeNow)
-							defer bm.Close(ctx)
+							defer bm.CloseShared(ctx)
 						case 1:
 							t.Logf("flushing")
 							bm.Flush(ctx)
@@ -1476,20 +1476,20 @@ func (s *contentManagerSuite) TestDeleteAndRecreate(t *testing.T) {
 			st := blobtesting.NewMapStorage(data, keyTime, faketime.Frozen(fakeTime))
 
 			bm := s.newTestContentManagerWithCustomTime(t, st, faketime.Frozen(fakeTime))
-			defer bm.Close(ctx)
+			defer bm.CloseShared(ctx)
 
 			content1 := writeContentAndVerify(ctx, t, bm, seededRandomData(10, 100))
 			bm.Flush(ctx)
 
 			// delete but at given timestamp but don't commit yet.
 			bm0 := s.newTestContentManagerWithCustomTime(t, st, faketime.AutoAdvance(tc.deletionTime, 1*time.Second))
-			defer bm0.Close(ctx)
+			defer bm0.CloseShared(ctx)
 
 			require.NoError(t, bm0.DeleteContent(ctx, content1))
 
 			// delete it at t0+10
 			bm1 := s.newTestContentManagerWithCustomTime(t, st, faketime.AutoAdvance(fakeTime.Add(10*time.Second), 1*time.Second))
-			defer bm1.Close(ctx)
+			defer bm1.CloseShared(ctx)
 
 			verifyContent(ctx, t, bm1, content1, seededRandomData(10, 100))
 			require.NoError(t, bm1.DeleteContent(ctx, content1))
@@ -1497,7 +1497,7 @@ func (s *contentManagerSuite) TestDeleteAndRecreate(t *testing.T) {
 
 			// recreate at t0+20
 			bm2 := s.newTestContentManagerWithCustomTime(t, st, faketime.AutoAdvance(fakeTime.Add(20*time.Second), 1*time.Second))
-			defer bm2.Close(ctx)
+			defer bm2.CloseShared(ctx)
 
 			content2 := writeContentAndVerify(ctx, t, bm2, seededRandomData(10, 100))
 			bm2.Flush(ctx)
@@ -1510,7 +1510,7 @@ func (s *contentManagerSuite) TestDeleteAndRecreate(t *testing.T) {
 			}
 
 			bm3 := s.newTestContentManager(t, st)
-			defer bm3.Close(ctx)
+			defer bm3.CloseShared(ctx)
 
 			dumpContentManagerData(t, data)
 			if tc.isVisible {
@@ -1791,7 +1791,7 @@ func (s *contentManagerSuite) TestContentWriteAliasing(t *testing.T) {
 	st := blobtesting.NewMapStorage(data, keyTime, faketime.Frozen(fakeTime))
 
 	bm := s.newTestContentManagerWithCustomTime(t, st, faketime.Frozen(fakeTime))
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 
 	contentData := []byte{100, 0, 0}
 	id1 := writeContentAndVerify(ctx, t, bm, contentData)
@@ -1843,7 +1843,7 @@ func (s *contentManagerSuite) TestContentReadAliasing(t *testing.T) {
 	st := blobtesting.NewMapStorage(data, keyTime, faketime.Frozen(fakeTime))
 
 	bm := s.newTestContentManagerWithCustomTime(t, st, faketime.Frozen(fakeTime))
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 
 	contentData := []byte{100, 0, 0}
 	id1 := writeContentAndVerify(ctx, t, bm, contentData)
@@ -1881,7 +1881,7 @@ func (s *contentManagerSuite) verifyVersionCompat(t *testing.T, writeVersion for
 	mgr := s.newTestContentManagerWithTweaks(t, st, &contentManagerTestTweaks{
 		formatVersion: writeVersion,
 	})
-	defer mgr.Close(ctx)
+	defer mgr.CloseShared(ctx)
 
 	dataSet := map[ID][]byte{}
 
@@ -1918,7 +1918,7 @@ func (s *contentManagerSuite) verifyVersionCompat(t *testing.T, writeVersion for
 
 	// create new manager that reads and writes using new version.
 	mgr = s.newTestContentManager(t, st)
-	defer mgr.Close(ctx)
+	defer mgr.CloseShared(ctx)
 
 	// make sure we can read everything
 	verifyContentManagerDataSet(ctx, t, mgr, dataSet)
@@ -1935,7 +1935,7 @@ func (s *contentManagerSuite) verifyVersionCompat(t *testing.T, writeVersion for
 
 	// now open one more manager
 	mgr = s.newTestContentManager(t, st)
-	defer mgr.Close(ctx)
+	defer mgr.CloseShared(ctx)
 	verifyContentManagerDataSet(ctx, t, mgr, dataSet)
 }
 
@@ -1993,13 +1993,13 @@ func (s *contentManagerSuite) verifyReadsOwnWrites(t *testing.T, st blob.Storage
 		if i%10 == 0 {
 			t.Logf("------- flushing & reopening -----")
 			require.NoError(t, bm.Flush(ctx))
-			require.NoError(t, bm.Close(ctx))
+			require.NoError(t, bm.CloseShared(ctx))
 			bm = s.newTestContentManagerWithTweaks(t, st, tweaks)
 		}
 	}
 
 	require.NoError(t, bm.Flush(ctx))
-	require.NoError(t, bm.Close(ctx))
+	require.NoError(t, bm.CloseShared(ctx))
 	bm = s.newTestContentManagerWithTweaks(t, st, tweaks)
 
 	for i := 0; i < len(ids); i++ {
@@ -2170,7 +2170,7 @@ func (s *contentManagerSuite) TestPrefetchContent(t *testing.T) {
 		maxPackSize: 20e6,
 	})
 
-	defer bm.Close(ctx)
+	defer bm.CloseShared(ctx)
 	bm.Flush(ctx)
 
 	// write 6 x 6 MB content in 2 blobs.
@@ -2372,7 +2372,7 @@ func (s *contentManagerSuite) newTestContentManagerWithTweaks(t *testing.T, st b
 	}
 
 	t.Cleanup(func() {
-		bm.Close(ctx)
+		bm.CloseShared(ctx)
 	})
 
 	bm.checkInvariantsOnUnlock = true

repo/grpc_repository_client.go

@@ -7,7 +7,6 @@ import (
 	"net/url"
 	"runtime"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/pkg/errors"
@@ -59,8 +58,8 @@ func errNoSessionResponse() error {
 // grpcRepositoryClient is an implementation of Repository that connects to an instance of
 // GPRC API server hosted by `kopia server`.
 type grpcRepositoryClient struct {
-	connRefCount *int32
-	conn         *grpc.ClientConn
+	*refCountedCloser
+	conn *grpc.ClientConn
 
 	innerSessionMutex sync.Mutex
 
@@ -461,11 +460,13 @@ func (r *grpcInnerSession) Flush(ctx context.Context) error {
 }
 
 func (r *grpcRepositoryClient) NewWriter(ctx context.Context, opt WriteSessionOptions) (context.Context, RepositoryWriter, error) {
-	w, err := newGRPCAPIRepositoryForConnection(ctx, r.conn, r.connRefCount, r.cliOpts, opt, r.contentCache, false, r.metricsRegistry)
+	w, err := newGRPCAPIRepositoryForConnection(ctx, r.conn, r.refCountedCloser, r.cliOpts, opt, r.contentCache, false, r.metricsRegistry)
 	if err != nil {
 		return nil, nil, err
 	}
 
+	w.addRef()
+
 	return ctx, w, nil
 }
 
@@ -727,25 +728,6 @@ func (r *grpcRepositoryClient) UpdateDescription(d string) {
 	r.cliOpts.Description = d
 }
 
-func (r *grpcRepositoryClient) Close(ctx context.Context) error {
-	if r.omgr == nil {
-		// already closed
-		return nil
-	}
-
-	r.omgr = nil
-
-	if atomic.AddInt32(r.connRefCount, -1) == 0 {
-		log(ctx).Debugf("closing GPRC connection to %v", r.conn.Target())
-
-		defer r.contentCache.Close(ctx)
-
-		return errors.Wrap(r.conn.Close(), "error closing GRPC connection")
-	}
-
-	return nil
-}
-
 var _ Repository = (*grpcRepositoryClient)(nil)
 
 type grpcCreds struct {
@@ -804,7 +786,18 @@ func OpenGRPCAPIRepository(ctx context.Context, si *APIServerInfo, cliOpts Clien
 		return nil, errors.Wrap(err, "dial error")
 	}
 
-	rep, err := newGRPCAPIRepositoryForConnection(ctx, conn, new(int32), cliOpts, WriteSessionOptions{}, contentCache, true, mr)
+	rcc := newRefCountedCloser(
+		func(ctx context.Context) error {
+			return errors.Wrap(conn.Close(), "error closing GRPC connection")
+		},
+		func(ctx context.Context) error {
+			contentCache.Close(ctx)
+			return nil
+		},
+		mr.Close,
+	)
+
+	rep, err := newGRPCAPIRepositoryForConnection(ctx, conn, rcc, cliOpts, WriteSessionOptions{}, contentCache, true, mr)
 	if err != nil {
 		return nil, err
 	}
@@ -880,7 +873,7 @@ func (r *grpcRepositoryClient) Metrics() *metrics.Registry {
 func newGRPCAPIRepositoryForConnection(
 	ctx context.Context,
 	conn *grpc.ClientConn,
-	connRefCount *int32,
+	sc *refCountedCloser,
 	cliOpts ClientOptions,
 	opt WriteSessionOptions,
 	contentCache *cache.PersistentCache,
@@ -892,7 +885,7 @@ func newGRPCAPIRepositoryForConnection(
 	}
 
 	rr := &grpcRepositoryClient{
-		connRefCount:       connRefCount,
+		refCountedCloser:   sc,
 		conn:               conn,
 		cliOpts:            cliOpts,
 		transparentRetries: transparentRetries,
@@ -923,8 +916,6 @@ func newGRPCAPIRepositoryForConnection(
 			return nil, errors.Wrap(err, "unable to initialize object manager")
 		}
 
-		atomic.AddInt32(connRefCount, 1)
-
 		return rr, nil
 	})
 }

repo/manifest/manifest_manager_test.go

@@ -167,7 +167,7 @@ func TestManifestInitCorruptedBlock(t *testing.T) {
 
 	bm0 := bm
 
-	t.Cleanup(func() { bm0.Close(ctx) })
+	t.Cleanup(func() { bm0.CloseShared(ctx) })
 
 	mgr, err := NewManager(ctx, bm, ManagerOptions{}, nil)
 	if err != nil {
@@ -193,7 +193,7 @@ func TestManifestInitCorruptedBlock(t *testing.T) {
 	bm, err = content.NewManagerForTesting(ctx, st, fop, nil, nil)
 	require.NoError(t, err)
 
-	t.Cleanup(func() { bm.Close(ctx) })
+	t.Cleanup(func() { bm.CloseShared(ctx) })
 
 	mgr, err = NewManager(ctx, bm, ManagerOptions{}, nil)
 	if err != nil {
@@ -325,7 +325,7 @@ func newManagerForTesting(ctx context.Context, t *testing.T, data blobtesting.Da
 	bm, err := content.NewManagerForTesting(ctx, st, fop, nil, nil)
 	require.NoError(t, err)
 
-	t.Cleanup(func() { bm.Close(ctx) })
+	t.Cleanup(func() { bm.CloseShared(ctx) })
 
 	mm, err := NewManager(ctx, bm, ManagerOptions{}, nil)
 	require.NoError(t, err)

repo/open.go

@@ -306,6 +306,11 @@ func openWithConfig(ctx context.Context, st blob.Storage, cliOpts ClientOptions,
 		return nil, errors.Wrap(ferr, "unable to open manifests")
 	}
 
+	rcc := newRefCountedCloser(
+		scm.CloseShared,
+		me.Close,
+	)
+
 	dr := &directRepository{
 		cmgr:  cm,
 		omgr:  om,
@@ -321,8 +326,9 @@ func openWithConfig(ctx context.Context, st blob.Storage, cliOpts ClientOptions,
 			nextWriterID:    new(int32),
 			throttler:       throttler,
 			metricsRegistry: me,
+
+			refCountedCloser: rcc,
 		},
-		closed: make(chan struct{}),
 	}
 
 	return dr, nil

repo/refcount_closer.go

@@ -0,0 +1,60 @@
+package repo
+
+import (
+	"context"
+	"sync/atomic"
+
+	"go.uber.org/multierr"
+)
+
+// closeFunc is a function to invoke when the last repository reference is closed.
+type closeFunc func(ctx context.Context) error
+
+type refCountedCloser struct {
+	refCount atomic.Int32
+	closed   atomic.Bool
+
+	closers []closeFunc
+}
+
+// Close decrements reference counter and invokes cleanup functions after last reference has been released.
+func (c *refCountedCloser) Close(ctx context.Context) error {
+	remaining := c.refCount.Add(-1)
+
+	if remaining != 0 {
+		return nil
+	}
+
+	if c.closed.Load() {
+		panic("already closed")
+	}
+
+	c.closed.Store(true)
+
+	var errors []error
+
+	for _, closer := range c.closers {
+		errors = append(errors, closer(ctx))
+	}
+
+	//nolint:wrapcheck
+	return multierr.Combine(errors...)
+}
+
+func (c *refCountedCloser) addRef() {
+	if c.closed.Load() {
+		panic("already closed")
+	}
+
+	c.refCount.Add(1)
+}
+
+func newRefCountedCloser(f ...closeFunc) *refCountedCloser {
+	rcc := &refCountedCloser{
+		closers: f,
+	}
+
+	rcc.addRef()
+
+	return rcc
+}

repo/repository.go

@@ -97,6 +97,8 @@ type directRepositoryParameters struct {
 	nextWriterID    *int32
 	throttler       throttling.SettableThrottler
 	metricsRegistry *metrics.Registry
+
+	*refCountedCloser
 }
 
 // directRepository is an implementation of repository that directly manipulates underlying storage.
@@ -108,8 +110,6 @@ type directRepository struct {
 	omgr  *object.Manager
 	mmgr  *manifest.Manager
 	sm    *content.SharedManager
-
-	closed chan struct{}
 }
 
 // DeriveKey derives encryption key of the provided length from the master key.
@@ -268,32 +268,13 @@ func (r *directRepository) NewDirectWriter(ctx context.Context, opt WriteSession
 		omgr:                       omgr,
 		mmgr:                       mmgr,
 		sm:                         r.sm,
-		closed:                     make(chan struct{}),
 	}
 
+	w.addRef()
+
 	return ctx, w, nil
 }
 
-// Close closes the repository and releases all resources.
-func (r *directRepository) Close(ctx context.Context) error {
-	select {
-	case <-r.closed:
-		// already closed
-		return nil
-
-	default:
-	}
-
-	// this will release shared manager and MAY release blob.Store (on last outstanding reference).
-	if err := r.cmgr.Close(ctx); err != nil {
-		return errors.Wrap(err, "error closing content-addressable storage manager")
-	}
-
-	close(r.closed)
-
-	return nil
-}
-
 // Flush waits for all in-flight writes to complete.
 func (r *directRepository) Flush(ctx context.Context) error {
 	if err := r.mmgr.Flush(ctx); err != nil {

tests/stress_test/stress_test.go

@@ -126,7 +126,7 @@ func stressWorker(ctx context.Context, t *testing.T, deadline time.Time, openMgr
 				return
 			}
 
-			if cerr := bm.Close(ctx); cerr != nil {
+			if cerr := bm.CloseShared(ctx); cerr != nil {
 				t.Errorf("close error: %v", cerr)
 				return
 			}