go: store/datas/pull: pull_chunk_tracker.go: Optimize memory use when backing up to an AWS remote.

PullChunkTracker is responsible for making the HasMany calls against the
destination and batching up absent hashes into HashSets which will be delivered
to GetManyCompressed and eventually written into table files which are
uploaded. This code is used for both pull and pull, when destination is "local"
database or when destination is the remote database respectively. It is used
when the remote is both doltremoteapi, thus every HasMany call is an RPC, and
when the remote is something like file:// or aws://, thus the table file
indexes for the remote are in memory and HasMany calls are very quick.o

Different operational characteristics of the various dependencies mean that
sometimes a Pull is prone to build up large sets of pending HasMany hashes,
whereas other times it is prone to build up large sets of absent hashes which
are waiting for the fetcher thread(s) to take them.

Previously, PullChunkTracker was structured to accumulate HasMany responses and
wait to batch them into appropriately-sized batches for GetManyCompressed until
the fetcher threads asked for them. This meant that if HasMany batches were
very small, because HasMany was very fast, we would accumulate are large number
of very small HashSets. These HashSets would take up large amounts of memory.
Accumulating the batches as the HasMany responses come in is more memory
efficient and should be no slower - we will always accumulate the full batches,
and in basically the same order.

Tested by pushing a large database to an AWS remote and memory profiling the
result.

Author: reltuk

go/store/datas/pull/pull_chunk_fetcher_test.go

@@ -183,3 +183,38 @@ var getManyerErr = fmt.Errorf("always return an error")
 func (errorGetManyer) GetManyCompressed(ctx context.Context, hashes hash.HashSet, found func(context.Context, nbs.ToChunker)) error {
 	return getManyerErr
 }
+
+func TestPop(t *testing.T) {
+	var backing [16]*int
+	for i := range 16 {
+		backing[i] = new(int)
+		*backing[i] = i
+	}
+	s := backing[:]
+	for i := range 16 {
+		assert.Len(t, s, 16-i)
+		var p *int
+		p, s = pop(s)
+		assert.Len(t, s, 16-i-1)
+		assert.Equal(t, i, *p)
+		assert.Nil(t, backing[16-i-1], "i is %d", i)
+	}
+}
+
+func TestAppendAbsent(t *testing.T) {
+	var absent []hash.HashSet
+	var hashes [16]hash.Hash
+	for i := range 16 {
+		hashes[i][0] = byte(i)
+	}
+	// Initial set is the full batch.
+	absent = appendAbsent(absent, hash.NewHashSet(hashes[:]...), 4)
+	assert.Len(t, absent, 1)
+	assert.Len(t, absent[0], 16)
+	// Next set get batched up.
+	absent = appendAbsent(absent, hash.NewHashSet(hashes[:]...), 4)
+	assert.Len(t, absent, 5)
+	for i := range 4 {
+		assert.Len(t, absent[i+1], 4)
+	}
+}

go/store/datas/pull/pull_chunk_tracker.go

@@ -131,7 +131,7 @@ func (t *PullChunkTracker) reqRespThread(ctx context.Context, initial hash.HashS
 
 	eg, ctx := errgroup.WithContext(ctx)
 
-	for i := 0; i < hasManyThreadCount; i++ {
+	for range hasManyThreadCount {
 		eg.Go(func() error {
 			return hasManyThread(ctx, t.cfg.HasManyer, hasManyReqCh, hasManyRespCh, doneCh)
 		})
@@ -169,61 +169,57 @@ func (t *PullChunkTracker) reqRespThread(ctx context.Context, initial hash.HashS
 
 			select {
 			case h, ok := <-t.uncheckedCh:
+				// uncheckedCh closing means the PullChunkTracker was closed.
 				if !ok {
 					return nil
 				}
+				// |h| is a hash we need to check against HasMany on the destination.
+				// It could become a hash we fetch from source and add to the destination
+				// or it could already be present.
 				if len(unchecked) == 0 || len(unchecked[len(unchecked)-1]) >= t.cfg.BatchSize {
 					outstanding += 1
 					unchecked = append(unchecked, make(hash.HashSet))
 				}
 				unchecked[len(unchecked)-1].Insert(h)
 			case resp := <-hasManyRespCh:
+				// A hasMany response came back.
 				outstanding -= 1
 				if resp.err != nil {
 					err = errors.Join(err, resp.err)
-				} else if len(resp.hs) > 0 {
-					absent = append(absent, resp.hs)
 				}
+				// Add all the resp.hs hashes, those which are not already present
+				// in dest, to our batches of absent hashes we will return through
+				// GetChunksToFetch.
+				absent = appendAbsent(absent, resp.hs, t.cfg.BatchSize)
 			case thisHasManyReqCh <- hasManyReq:
-				copy(unchecked[:], unchecked[1:])
-				if len(unchecked) > 1 {
-					unchecked[len(unchecked)-1] = nil
-				}
-				unchecked = unchecked[:len(unchecked)-1]
+				// We delivered a hasMany request to a hasManyThread.
+				// Remove it here. We do not need to update |outstanding|, since
+				// it was updated when we created the new batch in |unchecked|.
+				_, unchecked = pop(unchecked)
 			case <-t.processedCh:
+				// TickProcessed helps us keep track of chunks which we returned from GetChunksToFetch
+				// and which are still being processed by WalkAddrs. This only gets called after all
+				// unchecked hashes in the chunks have been delivered to us.
 				unprocessed -= 1
 			case req := <-thisReqCh:
+				// A request for GetChunksToFetch.
 				if err != nil {
+					// Deliver an error we experienced. HasMany can error, and this is where
+					// a client sees it, for example.
 					req.err = err
 					close(req.ready)
 					err = nil
 				} else if len(absent) == 0 {
+					// We have no more chunks to deliver. If len(absent) == 0 but we had
+					// unprocessed stuff, we would have had a |nil| |thisReqCh|. The fact that
+					// we accepted the request means we are ready to tell the client that there
+					// are no more chunks to fetch.
 					req.ok = false
 					close(req.ready)
 				} else {
+					// |absent[0]| is as full a batch as we have.
 					req.ok = true
-					req.hs = absent[0]
-					var i int
-					for i = 1; i < len(absent); i++ {
-						l := len(absent[i])
-						if len(req.hs)+l < t.cfg.BatchSize {
-							req.hs.InsertAll(absent[i])
-						} else {
-							for h := range absent[i] {
-								if len(req.hs) >= t.cfg.BatchSize {
-									break
-								}
-								req.hs.Insert(h)
-								absent[i].Remove(h)
-							}
-							break
-						}
-					}
-					copy(absent[:], absent[i:])
-					for j := len(absent) - i; j < len(absent); j++ {
-						absent[j] = nil
-					}
-					absent = absent[:len(absent)-i]
+					req.hs, absent = pop(absent)
 					unprocessed += len(req.hs)
 					close(req.ready)
 				}
@@ -236,6 +232,40 @@ func (t *PullChunkTracker) reqRespThread(ctx context.Context, initial hash.HashS
 	return eg.Wait()
 }
 
+// Pop returns the first element of s and the remaining elements of
+// s. It copies any suffix to the front of |s| and nils the last
+// element of |s| so that memory doesn't leak through |s[1:]| retaining
+// s[0].
+func pop[T any](s []T) (T, []T) {
+	ret := s[0]
+	copy(s[:], s[1:])
+	var empty T
+	s[len(s)-1] = empty
+	return ret, s[:len(s)-1]
+}
+
+// appendAbsent adds all elements in |toadd| to HashSets at the end of |absent|. It creates new HashSets
+// at the end of |absent| if inserting into an existing set would cause its size to exceed |sz|.
+//
+// As a special case, |appendAbsent| just returns []hash.HashSet{toadd} if |absent| is already empty.
+func appendAbsent(absent []hash.HashSet, toadd hash.HashSet, sz int) []hash.HashSet {
+	if len(toadd) == 0 {
+		return absent
+	}
+	// Don't bother splitting up toadd to |sz| if it is the first batch.
+	if len(absent) == 0 {
+		absent = append(absent, toadd)
+		return absent
+	}
+	for h := range toadd {
+		if len(absent[len(absent)-1]) >= sz {
+			absent = append(absent, make(hash.HashSet))
+		}
+		absent[len(absent)-1].Insert(h)
+	}
+	return absent
+}
+
 // Run by a PullChunkTracker, calls HasMany on a batch of addresses and delivers the results.
 func hasManyThread(ctx context.Context, hasManyer HasManyer, reqCh <-chan trackerHasManyReq, respCh chan<- trackerHasManyResp, doneCh <-chan struct{}) error {
 	for {