thread memory quota provider through tableSet

go/libraries/doltcore/dbfactory/aws.go

@@ -138,7 +138,8 @@ func (fact AWSFactory) newChunkStore(ctx context.Context, nbf *types.NomsBinForm
 	}
 
 	sess := session.Must(session.NewSessionWithOptions(opts))
-	return nbs.NewAWSStore(ctx, nbf.VersionString(), parts[0], dbName, parts[1], s3.New(sess), dynamodb.New(sess), defaultMemTableSize)
+	q := nbs.NewUnlimitedMemQuotaProvider()
+	return nbs.NewAWSStore(ctx, nbf.VersionString(), parts[0], dbName, parts[1], s3.New(sess), dynamodb.New(sess), defaultMemTableSize, q)
 }
 
 func validatePath(path string) (string, error) {

go/libraries/doltcore/dbfactory/file.go

@@ -57,8 +57,8 @@ func (fact FileFactory) CreateDB(ctx context.Context, nbf *types.NomsBinFormat,
 	if err != nil {
 		return nil, nil, err
 	}
-
-	newGenSt, err := nbs.NewLocalStore(ctx, nbf.VersionString(), path, defaultMemTableSize)
+	q := nbs.NewUnlimitedMemQuotaProvider()
+	newGenSt, err := nbs.NewLocalStore(ctx, nbf.VersionString(), path, defaultMemTableSize, q)
 
 	if err != nil {
 		return nil, nil, err
@@ -77,7 +77,7 @@ func (fact FileFactory) CreateDB(ctx context.Context, nbf *types.NomsBinFormat,
 		}
 	}
 
-	oldGenSt, err := nbs.NewLocalStore(ctx, newGenSt.Version(), oldgenPath, defaultMemTableSize)
+	oldGenSt, err := nbs.NewLocalStore(ctx, newGenSt.Version(), oldgenPath, defaultMemTableSize, q)
 
 	if err != nil {
 		return nil, nil, err

go/libraries/doltcore/dbfactory/gs.go

@@ -41,7 +41,8 @@ func (fact GSFactory) CreateDB(ctx context.Context, nbf *types.NomsBinFormat, ur
 	}
 
 	bs := blobstore.NewGCSBlobstore(gcs, urlObj.Host, urlObj.Path)
-	gcsStore, err := nbs.NewBSStore(ctx, nbf.VersionString(), bs, defaultMemTableSize)
+	q := nbs.NewUnlimitedMemQuotaProvider()
+	gcsStore, err := nbs.NewBSStore(ctx, nbf.VersionString(), bs, defaultMemTableSize, q)
 
 	if err != nil {
 		return nil, nil, err
@@ -67,7 +68,8 @@ func (fact LocalBSFactory) CreateDB(ctx context.Context, nbf *types.NomsBinForma
 	}
 
 	bs := blobstore.NewLocalBlobstore(absPath)
-	bsStore, err := nbs.NewBSStore(ctx, nbf.VersionString(), bs, defaultMemTableSize)
+	q := nbs.NewUnlimitedMemQuotaProvider()
+	bsStore, err := nbs.NewBSStore(ctx, nbf.VersionString(), bs, defaultMemTableSize, q)
 
 	if err != nil {
 		return nil, nil, err

go/performance/kvbench/prolly_store.go

@@ -38,7 +38,7 @@ func newMemoryProllyStore() keyValStore {
 func newNBSProllyStore(dir string) keyValStore {
 	ctx := context.Background()
 	verStr := types.Format_Default.VersionString()
-	cs, err := nbs.NewLocalStore(ctx, verStr, dir, defaultMemTableSize)
+	cs, err := nbs.NewLocalStore(ctx, verStr, dir, defaultMemTableSize, nbs.NewUnlimitedMemQuotaProvider())
 	if err != nil {
 		panic(err)
 	}

go/store/cmd/noms/noms_ds_test.go

@@ -45,7 +45,7 @@ type nomsDsTestSuite struct {
 func (s *nomsDsTestSuite) TestEmptyNomsDs() {
 	dir := s.DBDir
 
-	cs, err := nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), dir, clienttest.DefaultMemTableSize)
+	cs, err := nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), dir, clienttest.DefaultMemTableSize, nbs.NewUnlimitedMemQuotaProvider())
 	s.NoError(err)
 	ds := datas.NewDatabase(cs)
 
@@ -59,7 +59,7 @@ func (s *nomsDsTestSuite) TestEmptyNomsDs() {
 func (s *nomsDsTestSuite) TestNomsDs() {
 	dir := s.DBDir
 
-	cs, err := nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), dir, clienttest.DefaultMemTableSize)
+	cs, err := nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), dir, clienttest.DefaultMemTableSize, nbs.NewUnlimitedMemQuotaProvider())
 	s.NoError(err)
 	db := datas.NewDatabase(cs)
 

go/store/cmd/noms/noms_sync_test.go

@@ -44,7 +44,7 @@ type nomsSyncTestSuite struct {
 }
 
 func (s *nomsSyncTestSuite) TestSyncValidation() {
-	cs, err := nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), s.DBDir, clienttest.DefaultMemTableSize)
+	cs, err := nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), s.DBDir, clienttest.DefaultMemTableSize, nbs.NewUnlimitedMemQuotaProvider())
 	s.NoError(err)
 	sourceDB := datas.NewDatabase(cs)
 	source1, err := sourceDB.GetDataset(context.Background(), "src")
@@ -69,7 +69,7 @@ func (s *nomsSyncTestSuite) TestSyncValidation() {
 func (s *nomsSyncTestSuite) TestSync() {
 	defer s.NoError(file.RemoveAll(s.DBDir2))
 
-	cs, err := nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), s.DBDir, clienttest.DefaultMemTableSize)
+	cs, err := nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), s.DBDir, clienttest.DefaultMemTableSize, nbs.NewUnlimitedMemQuotaProvider())
 	s.NoError(err)
 	sourceDB := datas.NewDatabase(cs)
 	source1, err := sourceDB.GetDataset(context.Background(), "src")
@@ -88,7 +88,7 @@ func (s *nomsSyncTestSuite) TestSync() {
 	sout, _ := s.MustRun(main, []string{"sync", sourceSpec, sinkDatasetSpec})
 	s.Regexp("Synced", sout)
 
-	cs, err = nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), s.DBDir2, clienttest.DefaultMemTableSize)
+	cs, err = nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), s.DBDir2, clienttest.DefaultMemTableSize, nbs.NewUnlimitedMemQuotaProvider())
 	s.NoError(err)
 	db := datas.NewDatabase(cs)
 	dest, err := db.GetDataset(context.Background(), "dest")
@@ -101,7 +101,7 @@ func (s *nomsSyncTestSuite) TestSync() {
 	sout, _ = s.MustRun(main, []string{"sync", sourceDataset, sinkDatasetSpec})
 	s.Regexp("Synced", sout)
 
-	cs, err = nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), s.DBDir2, clienttest.DefaultMemTableSize)
+	cs, err = nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), s.DBDir2, clienttest.DefaultMemTableSize, nbs.NewUnlimitedMemQuotaProvider())
 	s.NoError(err)
 	db = datas.NewDatabase(cs)
 	dest, err = db.GetDataset(context.Background(), "dest")
@@ -118,7 +118,7 @@ func (s *nomsSyncTestSuite) TestSync() {
 	sout, _ = s.MustRun(main, []string{"sync", sourceDataset, sinkDatasetSpec})
 	s.Regexp("Created", sout)
 
-	cs, err = nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), s.DBDir2, clienttest.DefaultMemTableSize)
+	cs, err = nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), s.DBDir2, clienttest.DefaultMemTableSize, nbs.NewUnlimitedMemQuotaProvider())
 	s.NoError(err)
 	db = datas.NewDatabase(cs)
 	dest, err = db.GetDataset(context.Background(), "dest2")
@@ -130,7 +130,7 @@ func (s *nomsSyncTestSuite) TestSync() {
 func (s *nomsSyncTestSuite) TestSync_Issue2598() {
 	defer s.NoError(file.RemoveAll(s.DBDir2))
 
-	cs, err := nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), s.DBDir, clienttest.DefaultMemTableSize)
+	cs, err := nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), s.DBDir, clienttest.DefaultMemTableSize, nbs.NewUnlimitedMemQuotaProvider())
 	s.NoError(err)
 	sourceDB := datas.NewDatabase(cs)
 	// Create dataset "src1", which has a lineage of two commits.
@@ -153,7 +153,7 @@ func (s *nomsSyncTestSuite) TestSync_Issue2598() {
 	sourceDataset := spec.CreateValueSpecString("nbs", s.DBDir, "src1")
 	sinkDatasetSpec := spec.CreateValueSpecString("nbs", s.DBDir2, "dest")
 	sout, _ := s.MustRun(main, []string{"sync", sourceDataset, sinkDatasetSpec})
-	cs, err = nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), s.DBDir2, clienttest.DefaultMemTableSize)
+	cs, err = nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), s.DBDir2, clienttest.DefaultMemTableSize, nbs.NewUnlimitedMemQuotaProvider())
 	s.NoError(err)
 	db := datas.NewDatabase(cs)
 	dest, err := db.GetDataset(context.Background(), "dest")
@@ -165,7 +165,7 @@ func (s *nomsSyncTestSuite) TestSync_Issue2598() {
 	sourceDataset2 := spec.CreateValueSpecString("nbs", s.DBDir, "src2")
 	sinkDatasetSpec2 := spec.CreateValueSpecString("nbs", s.DBDir2, "dest2")
 	sout, _ = s.MustRun(main, []string{"sync", sourceDataset2, sinkDatasetSpec2})
-	cs, err = nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), s.DBDir2, clienttest.DefaultMemTableSize)
+	cs, err = nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), s.DBDir2, clienttest.DefaultMemTableSize, nbs.NewUnlimitedMemQuotaProvider())
 	s.NoError(err)
 	db = datas.NewDatabase(cs)
 	dest, err = db.GetDataset(context.Background(), "dest2")
@@ -179,7 +179,7 @@ func (s *nomsSyncTestSuite) TestSync_Issue2598() {
 
 func (s *nomsSyncTestSuite) TestRewind() {
 	var err error
-	cs, err := nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), s.DBDir, clienttest.DefaultMemTableSize)
+	cs, err := nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), s.DBDir, clienttest.DefaultMemTableSize, nbs.NewUnlimitedMemQuotaProvider())
 	s.NoError(err)
 	sourceDB := datas.NewDatabase(cs)
 	src, err := sourceDB.GetDataset(context.Background(), "foo")
@@ -195,7 +195,7 @@ func (s *nomsSyncTestSuite) TestRewind() {
 	sinkDatasetSpec := spec.CreateValueSpecString("nbs", s.DBDir, "foo")
 	s.MustRun(main, []string{"sync", sourceSpec, sinkDatasetSpec})
 
-	cs, err = nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), s.DBDir, clienttest.DefaultMemTableSize)
+	cs, err = nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), s.DBDir, clienttest.DefaultMemTableSize, nbs.NewUnlimitedMemQuotaProvider())
 	s.NoError(err)
 	db := datas.NewDatabase(cs)
 	dest, err := db.GetDataset(context.Background(), "foo")

go/store/datas/pull/puller_test.go

@@ -131,7 +131,7 @@ func tempDirDB(ctx context.Context) (types.ValueReadWriter, datas.Database, erro
 		return nil, nil, err
 	}
 
-	st, err := nbs.NewLocalStore(ctx, types.Format_Default.VersionString(), dir, clienttest.DefaultMemTableSize)
+	st, err := nbs.NewLocalStore(ctx, types.Format_Default.VersionString(), dir, clienttest.DefaultMemTableSize, nbs.NewUnlimitedMemQuotaProvider())
 	if err != nil {
 		return nil, nil, err
 	}

go/store/nbs/benchmarks/main.go

@@ -108,7 +108,7 @@ func main() {
 				d.PanicIfError(err)
 			}()
 			open = func() (chunks.ChunkStore, error) {
-				return nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), dir, bufSize)
+				return nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), dir, bufSize, nbs.NewUnlimitedMemQuotaProvider())
 			}
 			reset = func() {
 				err := file.RemoveAll(dir)
@@ -138,7 +138,7 @@ func main() {
 		} else if *toAWS != "" {
 			sess := session.Must(session.NewSession(aws.NewConfig().WithRegion("us-west-2")))
 			open = func() (chunks.ChunkStore, error) {
-				return nbs.NewAWSStore(context.Background(), types.Format_Default.VersionString(), dynamoTable, *toAWS, s3Bucket, s3.New(sess), dynamodb.New(sess), bufSize)
+				return nbs.NewAWSStore(context.Background(), types.Format_Default.VersionString(), dynamoTable, *toAWS, s3Bucket, s3.New(sess), dynamodb.New(sess), bufSize, nbs.NewUnlimitedMemQuotaProvider())
 			}
 			reset = func() {
 				ddb := dynamodb.New(sess)
@@ -160,12 +160,12 @@ func main() {
 	} else {
 		if *useNBS != "" {
 			open = func() (chunks.ChunkStore, error) {
-				return nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), *useNBS, bufSize)
+				return nbs.NewLocalStore(context.Background(), types.Format_Default.VersionString(), *useNBS, bufSize, nbs.NewUnlimitedMemQuotaProvider())
 			}
 		} else if *useAWS != "" {
 			sess := session.Must(session.NewSession(aws.NewConfig().WithRegion("us-west-2")))
 			open = func() (chunks.ChunkStore, error) {
-				return nbs.NewAWSStore(context.Background(), types.Format_Default.VersionString(), dynamoTable, *useAWS, s3Bucket, s3.New(sess), dynamodb.New(sess), bufSize)
+				return nbs.NewAWSStore(context.Background(), types.Format_Default.VersionString(), dynamoTable, *useAWS, s3Bucket, s3.New(sess), dynamodb.New(sess), bufSize, nbs.NewUnlimitedMemQuotaProvider())
 			}
 		}
 		writeDB = func() {}

go/store/nbs/block_store_test.go

@@ -60,7 +60,7 @@ func (suite *BlockStoreSuite) SetupTest() {
 	var err error
 	suite.dir, err = os.MkdirTemp("", "")
 	suite.NoError(err)
-	suite.store, err = NewLocalStore(context.Background(), constants.FormatDefaultString, suite.dir, testMemTableSize)
+	suite.store, err = NewLocalStore(context.Background(), constants.FormatDefaultString, suite.dir, testMemTableSize, NewUnlimitedMemQuotaProvider())
 	suite.NoError(err)
 	suite.putCountFn = func() int {
 		return int(suite.store.putCount)
@@ -78,7 +78,7 @@ func (suite *BlockStoreSuite) TearDownTest() {
 
 func (suite *BlockStoreSuite) TestChunkStoreMissingDir() {
 	newDir := filepath.Join(suite.dir, "does-not-exist")
-	_, err := NewLocalStore(context.Background(), constants.FormatDefaultString, newDir, testMemTableSize)
+	_, err := NewLocalStore(context.Background(), constants.FormatDefaultString, newDir, testMemTableSize, NewUnlimitedMemQuotaProvider())
 	suite.Error(err)
 }
 
@@ -87,7 +87,7 @@ func (suite *BlockStoreSuite) TestChunkStoreNotDir() {
 	_, err := os.Create(existingFile)
 	suite.NoError(err)
 
-	_, err = NewLocalStore(context.Background(), constants.FormatDefaultString, existingFile, testMemTableSize)
+	_, err = NewLocalStore(context.Background(), constants.FormatDefaultString, existingFile, testMemTableSize, NewUnlimitedMemQuotaProvider())
 	suite.Error(err)
 }
 
@@ -276,7 +276,7 @@ func (suite *BlockStoreSuite) TestChunkStoreFlushOptimisticLockFail() {
 	root, err := suite.store.Root(context.Background())
 	suite.NoError(err)
 
-	interloper, err := NewLocalStore(context.Background(), constants.FormatDefaultString, suite.dir, testMemTableSize)
+	interloper, err := NewLocalStore(context.Background(), constants.FormatDefaultString, suite.dir, testMemTableSize, NewUnlimitedMemQuotaProvider())
 	suite.NoError(err)
 	err = interloper.Put(context.Background(), c1)
 	suite.NoError(err)
@@ -322,7 +322,7 @@ func (suite *BlockStoreSuite) TestChunkStoreRebaseOnNoOpFlush() {
 	input1 := []byte("abc")
 	c1 := chunks.NewChunk(input1)
 
-	interloper, err := NewLocalStore(context.Background(), constants.FormatDefaultString, suite.dir, testMemTableSize)
+	interloper, err := NewLocalStore(context.Background(), constants.FormatDefaultString, suite.dir, testMemTableSize, NewUnlimitedMemQuotaProvider())
 	suite.NoError(err)
 	err = interloper.Put(context.Background(), c1)
 	suite.NoError(err)
@@ -358,7 +358,7 @@ func (suite *BlockStoreSuite) TestChunkStorePutWithRebase() {
 	root, err := suite.store.Root(context.Background())
 	suite.NoError(err)
 
-	interloper, err := NewLocalStore(context.Background(), constants.FormatDefaultString, suite.dir, testMemTableSize)
+	interloper, err := NewLocalStore(context.Background(), constants.FormatDefaultString, suite.dir, testMemTableSize, NewUnlimitedMemQuotaProvider())
 	suite.NoError(err)
 	err = interloper.Put(context.Background(), c1)
 	suite.NoError(err)
@@ -443,10 +443,15 @@ func TestBlockStoreConjoinOnCommit(t *testing.T) {
 	t.Run("NoConjoin", func(t *testing.T) {
 		mm := makeManifestManager(&fakeManifest{})
 		p := newFakeTablePersister()
+		q := NewUnlimitedMemQuotaProvider()
+		defer func() {
+			require.EqualValues(t, 0, q.Usage())
+		}()
 		c := &fakeConjoiner{}
 
-		smallTableStore, err := newNomsBlockStore(context.Background(), constants.FormatDefaultString, mm, p, c, testMemTableSize)
+		smallTableStore, err := newNomsBlockStore(context.Background(), constants.FormatDefaultString, mm, p, q, c, testMemTableSize)
 		require.NoError(t, err)
+		defer smallTableStore.Close()
 
 		root, err := smallTableStore.Root(context.Background())
 		require.NoError(t, err)
@@ -477,7 +482,10 @@ func TestBlockStoreConjoinOnCommit(t *testing.T) {
 	t.Run("ConjoinSuccess", func(t *testing.T) {
 		fm := &fakeManifest{}
 		p := newFakeTablePersister()
-
+		q := NewUnlimitedMemQuotaProvider()
+		defer func() {
+			require.EqualValues(t, 0, q.Usage())
+		}()
 		srcs := makeTestSrcs(t, []uint32{1, 1, 3, 7}, p)
 		upstream, err := toSpecs(srcs)
 		require.NoError(t, err)
@@ -486,8 +494,9 @@ func TestBlockStoreConjoinOnCommit(t *testing.T) {
 			[]cannedConjoin{makeCanned(upstream[:2], upstream[2:], p)},
 		}
 
-		smallTableStore, err := newNomsBlockStore(context.Background(), constants.FormatDefaultString, makeManifestManager(fm), p, c, testMemTableSize)
+		smallTableStore, err := newNomsBlockStore(context.Background(), constants.FormatDefaultString, makeManifestManager(fm), p, q, c, testMemTableSize)
 		require.NoError(t, err)
+		defer smallTableStore.Close()
 
 		root, err := smallTableStore.Root(context.Background())
 		require.NoError(t, err)
@@ -509,7 +518,10 @@ func TestBlockStoreConjoinOnCommit(t *testing.T) {
 	t.Run("ConjoinRetry", func(t *testing.T) {
 		fm := &fakeManifest{}
 		p := newFakeTablePersister()
-
+		q := NewUnlimitedMemQuotaProvider()
+		defer func() {
+			require.EqualValues(t, 0, q.Usage())
+		}()
 		srcs := makeTestSrcs(t, []uint32{1, 1, 3, 7, 13}, p)
 		upstream, err := toSpecs(srcs)
 		require.NoError(t, err)
@@ -521,8 +533,9 @@ func TestBlockStoreConjoinOnCommit(t *testing.T) {
 			},
 		}
 
-		smallTableStore, err := newNomsBlockStore(context.Background(), constants.FormatDefaultString, makeManifestManager(fm), p, c, testMemTableSize)
+		smallTableStore, err := newNomsBlockStore(context.Background(), constants.FormatDefaultString, makeManifestManager(fm), p, q, c, testMemTableSize)
 		require.NoError(t, err)
+		defer smallTableStore.Close()
 
 		root, err := smallTableStore.Root(context.Background())
 		require.NoError(t, err)

go/store/nbs/cache.go

@@ -43,7 +43,7 @@ func NewCache(ctx context.Context) (*NomsBlockCache, error) {
 		return nil, err
 	}
 
-	store, err := NewLocalStore(ctx, types.Format_Default.VersionString(), dir, defaultCacheMemTableSize)
+	store, err := NewLocalStore(ctx, types.Format_Default.VersionString(), dir, defaultCacheMemTableSize, NewUnlimitedMemQuotaProvider())
 
 	if err != nil {
 		return nil, err

go/store/nbs/frag/main.go

@@ -70,15 +70,15 @@ func main() {
 	var store *nbs.NomsBlockStore
 	if *dir != "" {
 		var err error
-		store, err = nbs.NewLocalStore(ctx, types.Format_Default.VersionString(), *dir, memTableSize)
+		store, err = nbs.NewLocalStore(ctx, types.Format_Default.VersionString(), *dir, memTableSize, nbs.NewUnlimitedMemQuotaProvider())
 		d.PanicIfError(err)
 
 		*dbName = *dir
 	} else if *table != "" && *bucket != "" && *dbName != "" {
 		sess := session.Must(session.NewSession(aws.NewConfig().WithRegion("us-west-2")))
 
 		var err error
-		store, err = nbs.NewAWSStore(context.Background(), types.Format_Default.VersionString(), *table, *dbName, *bucket, s3.New(sess), dynamodb.New(sess), memTableSize)
+		store, err = nbs.NewAWSStore(context.Background(), types.Format_Default.VersionString(), *table, *dbName, *bucket, s3.New(sess), dynamodb.New(sess), memTableSize, nbs.NewUnlimitedMemQuotaProvider())
 		d.PanicIfError(err)
 	} else {
 		log.Fatalf("Must set either --dir or ALL of --table, --bucket and --db\n")

go/store/nbs/manifest.go

@@ -450,6 +450,11 @@ func (ts tableSpec) GetChunkCount() uint32 {
 	return ts.chunkCount
 }
 
+func (ts tableSpec) GetMemorySize() uint64 {
+	n := ts.GetChunkCount()
+	return memSize(n)
+}
+
 func tableSpecsToMap(specs []tableSpec) map[string]int {
 	m := make(map[string]int)
 	for _, spec := range specs {