Refactor sequence processing for parallel loading and memory management

This commit introduces significant changes to the sequence processing pipeline, implementing parallel loading, k-mer budget control, explicit memory management, and refined logic for pair evaluation and resolution.

pkg/obikmer/seqquorum.go

@@ -24,9 +24,12 @@ import (
 //  2. CountAtLeast processes the partitions in parallel: each one is loaded,
 //     expanded to (canonical k-mer, sequence id) pairs, sorted, and reduced.
 //
-// Memory during phase 2 is about 16 bytes per k-mer occurrence of the largest
-// partitions being processed at the same time, so P must grow with the total
-// input size (total bases / P is the occurrence count of an average partition).
+// Memory during phase 2 is 16 bytes per k-mer occurrence of the partitions
+// being processed at the same time. Their number is bounded by an occurrence
+// budget (see SetOccurrenceBudget), so that the memory does not depend on how
+// the k-mers happen to be spread over the partitions; P only has to be large
+// enough for the largest partition to be small against the budget (total bases
+// / P is the occurrence count of an average partition, see DefaultPartitions).
 //
 // Bases other than A, C, G, T, U break a sequence into independent runs: no
 // k-mer ever contains an ambiguous base.
@@ -39,13 +42,16 @@ type SeqKmerCounter struct {
 	mu       sync.Mutex
 	firstErr error
 	sealed   bool
+
+	budget int // k-mer occurrences held at once in phase 2
 }
 
 type seqPartWriter struct {
-	mu   sync.Mutex
-	w    *bufio.Writer
-	file *os.File
-	path string
+	mu    sync.Mutex
+	w     *bufio.Writer
+	file  *os.File
+	path  string
+	kmers int // k-mer occurrences written, so that the partition is loaded in an array of the exact size
 }
 
 // SeqKmerCount is the result for one k-mer.
@@ -78,7 +84,7 @@ func NewSeqKmerCounter(parentDir string, k, m, partitions int) (*SeqKmerCounter,
 		return nil, err
 	}
 
-	c := &SeqKmerCounter{dir: dir, k: k, m: m, p: partitions}
+	c := &SeqKmerCounter{dir: dir, k: k, m: m, p: partitions, budget: DefaultOccurrenceBudget}
 	c.parts = make([]*seqPartWriter, partitions)
 	for i := range c.parts {
 		path := filepath.Join(dir, fmt.Sprintf("part_%04d.sqk", i))
@@ -92,6 +98,38 @@ func NewSeqKmerCounter(parentDir string, k, m, partitions int) (*SeqKmerCounter,
 	return c, nil
 }
 
+// DefaultOccurrenceBudget is the number of k-mer occurrences phase 2 holds at
+// once by default: 32 million, 512 MB. A partition larger than the budget is
+// processed alone. Measured on 16 cores, 526 million occurrences to count: with
+// 32M the count is as fast as without any limit; 16M make it 15% slower, 8M
+// 60% slower.
+const DefaultOccurrenceBudget = 32 << 20
+
+// SetOccurrenceBudget bounds the k-mer occurrences (16 bytes each) that
+// CountAtLeast holds in memory at once. It must be called before CountAtLeast.
+func (c *SeqKmerCounter) SetOccurrenceBudget(occurrences int) {
+	if occurrences > 0 {
+		c.budget = occurrences
+	}
+}
+
+// partitionOccurrences is the occurrence count aimed at for an average
+// partition by DefaultPartitions (a partition of 1M occurrences is 16 MB).
+const partitionOccurrences = 1 << 20
+
+// DefaultPartitions returns a partition count suited to a total of kmers
+// k-mer occurrences: 256 at least, and more, as a power of two up to 1024, when
+// the average partition would otherwise exceed about a million occurrences.
+// The largest partitions hold several times the average. The limit keeps the
+// number of simultaneously open files reasonable.
+func DefaultPartitions(kmers int) int {
+	p := 256
+	for p < 1024 && kmers/p > partitionOccurrences {
+		p *= 2
+	}
+	return p
+}
+
 // K returns the k-mer size.
 func (c *SeqKmerCounter) K() int { return c.k }
 
@@ -148,6 +186,7 @@ func (c *SeqKmerCounter) addRun(id uint32, run []byte, scratch []byte) []byte {
 		pw := c.parts[int(sk.Minimizer%uint64(c.p))]
 		pw.mu.Lock()
 		_, err := pw.w.Write(scratch)
+		pw.kmers += len(sk.Sequence) - c.k + 1
 		pw.mu.Unlock()
 		if err != nil {
 			c.setErr(err)

pkg/obikmer/seqquorum_budget_test.go

@@ -0,0 +1,118 @@
+package obikmer
+
+import (
+	"math/rand"
+	"reflect"
+	"sync"
+	"testing"
+	"time"
+)
+
+// The occurrence budget only bounds the memory: any budget, down to one
+// occurrence (every partition then goes alone), gives the same counts.
+func TestSeqKmerCounterBudgetDoesNotChangeTheCounts(t *testing.T) {
+	rng := rand.New(rand.NewSource(7))
+	seqs := randomSeqs(rng, 40, 20, 500, true)
+	build := func(budget int) []SeqKmerCount {
+		c, err := NewSeqKmerCounter(t.TempDir(), 9, 4, 11)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer c.Close()
+		if budget > 0 {
+			c.SetOccurrenceBudget(budget)
+		}
+		for i, s := range seqs {
+			c.AddSequence(uint32(i), s, i%2 == 0)
+		}
+		total := 0
+		for _, pw := range c.parts {
+			total += pw.kmers
+		}
+		if total == 0 {
+			t.Fatal("no occurrence counted")
+		}
+		got, err := c.CountAtLeast(2)
+		if err != nil {
+			t.Fatal(err)
+		}
+		return got
+	}
+	want := build(0)
+	if len(want) == 0 {
+		t.Fatal("the test needs k-mers shared by several sequences")
+	}
+	for _, budget := range []int{1, 100, 5000} {
+		if got := build(budget); !reflect.DeepEqual(got, want) {
+			t.Errorf("budget %d changes the counts", budget)
+		}
+	}
+}
+
+// The occurrences announced for a partition are those that its file holds.
+func TestSeqKmerCounterPartitionOccurrenceCount(t *testing.T) {
+	rng := rand.New(rand.NewSource(3))
+	seqs := randomSeqs(rng, 25, 10, 300, true)
+	c, err := NewSeqKmerCounter(t.TempDir(), 8, 3, 5)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer c.Close()
+	for i, s := range seqs {
+		c.AddSequence(uint32(i), s, i%3 == 0)
+	}
+	if err := c.seal(); err != nil {
+		t.Fatal(err)
+	}
+	for p, pw := range c.parts {
+		occ, err := c.loadPartition(p, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(occ) != pw.kmers {
+			t.Errorf("partition %d: %d occurrences loaded, %d counted when written", p, len(occ), pw.kmers)
+		}
+	}
+}
+
+func TestOccurrenceBudget(t *testing.T) {
+	b := newOccurrenceBudget(10)
+	if got := b.acquire(4); got != 4 {
+		t.Fatalf("acquire(4) = %d", got)
+	}
+	// 7 does not fit in the 6 left: it must wait for the release
+	done := make(chan int)
+	go func() { done <- b.acquire(7) }()
+	select {
+	case <-done:
+		t.Fatal("acquire(7) must wait")
+	case <-time.After(50 * time.Millisecond):
+	}
+	b.release(4)
+	if got := <-done; got != 7 {
+		t.Fatalf("acquire(7) = %d", got)
+	}
+	b.release(7)
+
+	// a request above the whole budget takes all of it, once everything is free
+	if got := b.acquire(1000); got != 10 {
+		t.Fatalf("an oversized request must take the whole budget, got %d", got)
+	}
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() { defer wg.Done(); b.release(b.acquire(3)) }()
+	time.Sleep(20 * time.Millisecond)
+	b.release(10)
+	wg.Wait()
+}
+
+func TestDefaultPartitions(t *testing.T) {
+	for _, tc := range []struct{ kmers, want int }{
+		{0, 256}, {5_000_000, 256}, {256 << 20, 256}, {257 << 20, 512}, {526_000_000, 512},
+		{600 << 20, 1024}, {1 << 40, 1024},
+	} {
+		if got := DefaultPartitions(tc.kmers); got != tc.want {
+			t.Errorf("DefaultPartitions(%d) = %d, want %d", tc.kmers, got, tc.want)
+		}
+	}
+}

pkg/obikmer/seqquorum_count.go

@@ -41,6 +41,16 @@ func (c *SeqKmerCounter) CountAtLeast(minSeq int) ([]SeqKmerCount, error) {
 	}
 
 	workers := min(obidefault.ParallelWorkers(), c.p)
+
+	// the largest partitions first: they are the hardest to fit in the budget
+	// and the last ones to finish, if they come late
+	order := make([]int, c.p)
+	for i := range order {
+		order[i] = i
+	}
+	slices.SortStableFunc(order, func(a, b int) int { return cmp.Compare(c.parts[b].kmers, c.parts[a].kmers) })
+	budget := newOccurrenceBudget(c.budget)
+
 	jobs := make(chan int)
 	results := make([][]SeqKmerCount, c.p)
 
@@ -52,13 +62,19 @@ func (c *SeqKmerCounter) CountAtLeast(minSeq int) ([]SeqKmerCount, error) {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			var occ []kmerOcc // reused between partitions
+			var occ []kmerOcc // reused between partitions of similar sizes
 			for p := range jobs {
+				n := c.parts[p].kmers
+				held := budget.acquire(n)
+				if cap(occ) < n || cap(occ) > 2*n+1024 {
+					occ = make([]kmerOcc, 0, n)
+				}
 				var err error
 				occ, err = c.loadPartition(p, occ[:0])
 				if err == nil {
 					results[p] = reduceOccurrences(occ, minSeq)
 				}
+				budget.release(held)
 				if err != nil {
 					errMu.Lock()
 					if firstErr == nil {
@@ -69,7 +85,7 @@ func (c *SeqKmerCounter) CountAtLeast(minSeq int) ([]SeqKmerCount, error) {
 			}
 		}()
 	}
-	for p := range c.p {
+	for _, p := range order {
 		jobs <- p
 	}
 	close(jobs)
@@ -153,3 +169,39 @@ func reduceOccurrences(occ []kmerOcc, minSeq int) []SeqKmerCount {
 	}
 	return out
 }
+
+// occurrenceBudget bounds the k-mer occurrences held in memory at once by the
+// goroutines that load the partitions.
+type occurrenceBudget struct {
+	mu    sync.Mutex
+	freed *sync.Cond
+	total int
+	free  int
+}
+
+func newOccurrenceBudget(total int) *occurrenceBudget {
+	b := &occurrenceBudget{total: total, free: total}
+	b.freed = sync.NewCond(&b.mu)
+	return b
+}
+
+// acquire waits until n occurrences fit in the budget and takes them; a request
+// above the whole budget waits for the budget to be free and takes all of it.
+// It returns what was taken, to give back with release.
+func (b *occurrenceBudget) acquire(n int) int {
+	n = min(n, b.total)
+	b.mu.Lock()
+	for b.free < n {
+		b.freed.Wait()
+	}
+	b.free -= n
+	b.mu.Unlock()
+	return n
+}
+
+func (b *occurrenceBudget) release(n int) {
+	b.mu.Lock()
+	b.free += n
+	b.mu.Unlock()
+	b.freed.Broadcast()
+}