fix more stats races

go/libraries/doltcore/sqle/enginetest/dolt_engine_test.go

@@ -1954,6 +1954,7 @@ func TestStatsAutoRefreshConcurrency(t *testing.T) {
 	// create engine
 	harness := newDoltHarness(t)
 	harness.Setup(setup.MydbData)
+	harness.configureStats = true
 	engine := mustNewEngine(t, harness)
 	defer engine.Close()
 
@@ -1973,8 +1974,9 @@ func TestStatsAutoRefreshConcurrency(t *testing.T) {
 	// it is important to use new sessions for this test, to avoid working root conflicts
 	readCtx := enginetest.NewSession(harness)
 	writeCtx := enginetest.NewSession(harness)
+	refreshCtx := enginetest.NewSession(harness)
 	newCtx := func(context.Context) (*sql.Context, error) {
-		return enginetest.NewSession(harness), nil
+		return refreshCtx, nil
 	}
 
 	err := statsProv.InitAutoRefreshWithParams(newCtx, sqlDb.Name(), bThreads, intervalSec, thresholdf64, branches)
@@ -1988,7 +1990,7 @@ func TestStatsAutoRefreshConcurrency(t *testing.T) {
 		require.NoError(t, err)
 	}
 
-	iters := 1_000
+	iters := 50
 	{
 		// 3 threads to test auto-refresh/DML concurrency safety
 		// - auto refresh (read + write)

go/libraries/doltcore/sqle/statsnoms/database.go

@@ -151,9 +151,6 @@ func (n *NomsStatsDatabase) LoadBranchStats(ctx *sql.Context, branch string) err
 }
 
 func (n *NomsStatsDatabase) getBranchStats(branch string) dbStats {
-	n.mu.Lock()
-	defer n.mu.Unlock()
-
 	for i, b := range n.branches {
 		if strings.EqualFold(b, branch) {
 			return n.stats[i]
@@ -163,12 +160,16 @@ func (n *NomsStatsDatabase) getBranchStats(branch string) dbStats {
 }
 
 func (n *NomsStatsDatabase) GetStat(branch string, qual sql.StatQualifier) (*statspro.DoltStats, bool) {
+	n.mu.Lock()
+	defer n.mu.Unlock()
 	stats := n.getBranchStats(branch)
 	ret, ok := stats[qual]
 	return ret, ok
 }
 
 func (n *NomsStatsDatabase) ListStatQuals(branch string) []sql.StatQualifier {
+	n.mu.Lock()
+	defer n.mu.Unlock()
 	stats := n.getBranchStats(branch)
 	var ret []sql.StatQualifier
 	for qual, _ := range stats {
@@ -286,7 +287,11 @@ func (n *NomsStatsDatabase) ReplaceChunks(ctx context.Context, branch string, qu
 		if err != nil {
 			return err
 		}
-		dbStat[qual].Hist = targetBuckets
+		newStat, err := dbStat[qual].WithHistogram(targetBuckets)
+		if err != nil {
+			return err
+		}
+		dbStat[qual] = newStat.(*statspro.DoltStats)
 	} else {
 		dbStat[qual] = statspro.NewDoltStats()
 	}

go/libraries/doltcore/sqle/statspro/dolt_stats.go

@@ -161,6 +161,8 @@ func (s *DoltStats) ToInterface() (interface{}, error) {
 }
 
 func (s *DoltStats) WithHistogram(h sql.Histogram) (sql.Statistic, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
 	ret := *s
 	ret.Hist = nil
 	for _, b := range h {
@@ -174,6 +176,8 @@ func (s *DoltStats) WithHistogram(h sql.Histogram) (sql.Statistic, error) {
 }
 
 func (s *DoltStats) Histogram() sql.Histogram {
+	s.mu.Lock()
+	defer s.mu.Unlock()
 	return s.Hist
 }