X-Git-Url: https://git.arvados.org/lightning.git/blobdiff_plain/df2c5ad435cc3f26f01a077fd3666dceee1e3793..e4839e500e2da73b476f095425c0347ae7c5de97:/tilelib.go

diff --git a/tilelib.go b/tilelib.go
index 8960d0e86d..550563342c 100644
--- a/tilelib.go
+++ b/tilelib.go
@@ -1,3 +1,7 @@
+// Copyright (C) The Lightning Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
 package lightning
 
 import (
@@ -62,11 +66,17 @@ type tileLibrary struct {
 	variant        [][][blake2b.Size256]byte
 	refseqs        map[string]map[string][]tileLibRef
 	compactGenomes map[string][]tileVariantID
-	// count [][]int
-	seq      map[[blake2b.Size256]byte][]byte
-	variants int64
+	seq2           map[[2]byte]map[[blake2b.Size256]byte][]byte
+	seq2lock       map[[2]byte]sync.Locker
+	variants       int64
 	// if non-nil, write out any tile variants added while tiling
 	encoder *gob.Encoder
+	// set Ref flag when writing new variants to encoder
+	encodeRef bool
+
+	onAddTileVariant func(libref tileLibRef, hash [blake2b.Size256]byte, seq []byte) error
+	onAddGenome      func(CompactGenome) error
+	onAddRefseq      func(CompactSequence) error
 
 	mtx   sync.RWMutex
 	vlock []sync.Locker
@@ -112,12 +122,12 @@ func (tilelib *tileLibrary) loadTileVariants(tvs []TileVariant, variantmap map[t
 	for _, tv := range tvs {
 		// Assign a new variant ID (unique across all inputs)
 		// for each input variant.
-		variantmap[tileLibRef{Tag: tv.Tag, Variant: tv.Variant}] = tilelib.getRef(tv.Tag, tv.Sequence).Variant
+		variantmap[tileLibRef{Tag: tv.Tag, Variant: tv.Variant}] = tilelib.getRef(tv.Tag, tv.Sequence, tv.Ref).Variant
 	}
 	return nil
 }
 
-func (tilelib *tileLibrary) loadCompactGenomes(cgs []CompactGenome, variantmap map[tileLibRef]tileVariantID, onLoadGenome func(CompactGenome)) error {
+func (tilelib *tileLibrary) loadCompactGenomes(cgs []CompactGenome, variantmap map[tileLibRef]tileVariantID) error {
 	log.Debugf("loadCompactGenomes: %d", len(cgs))
 	var wg sync.WaitGroup
 	errs := make(chan error, 1)
@@ -143,11 +153,18 @@ func (tilelib *tileLibrary) loadCompactGenomes(cgs []CompactGenome, variantmap m
 					}
 					return
 				}
-				log.Tracef("loadCompactGenomes: cg %s tag %d variant %d => %d", cg.Name, tag, variant, newvariant)
+				// log.Tracef("loadCompactGenomes: cg %s tag %d variant %d => %d", cg.Name, tag, variant, newvariant)
 				cg.Variants[i] = newvariant
 			}
-			if onLoadGenome != nil {
-				onLoadGenome(cg)
+			if tilelib.onAddGenome != nil {
+				err := tilelib.onAddGenome(cg)
+				if err != nil {
+					select {
+					case errs <- err:
+					default:
+					}
+					return
+				}
 			}
 			if tilelib.encoder != nil {
 				err := tilelib.encoder.Encode(LibraryEntry{
@@ -174,8 +191,9 @@ func (tilelib *tileLibrary) loadCompactGenomes(cgs []CompactGenome, variantmap m
 }
 
 func (tilelib *tileLibrary) loadCompactSequences(cseqs []CompactSequence, variantmap map[tileLibRef]tileVariantID) error {
-	log.Debugf("loadCompactSequences: %d", len(cseqs))
+	log.Infof("loadCompactSequences: %d todo", len(cseqs))
 	for _, cseq := range cseqs {
+		log.Infof("loadCompactSequences: checking %s", cseq.Name)
 		for _, tseq := range cseq.TileSequences {
 			for i, libref := range tseq {
 				if libref.Variant == 0 {
@@ -198,6 +216,13 @@ func (tilelib *tileLibrary) loadCompactSequences(cseqs []CompactSequence, varian
 				return err
 			}
 		}
+		if tilelib.onAddRefseq != nil {
+			err := tilelib.onAddRefseq(cseq)
+			if err != nil {
+				return err
+			}
+		}
+		log.Infof("loadCompactSequences: checking %s done", cseq.Name)
 	}
 	tilelib.mtx.Lock()
 	defer tilelib.mtx.Unlock()
@@ -207,56 +232,57 @@ func (tilelib *tileLibrary) loadCompactSequences(cseqs []CompactSequence, varian
 	for _, cseq := range cseqs {
 		tilelib.refseqs[cseq.Name] = cseq.TileSequences
 	}
+	log.Info("loadCompactSequences: done")
 	return nil
 }
 
-func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string, onLoadGenome func(CompactGenome)) error {
+func allFiles(path string, re *regexp.Regexp) ([]string, error) {
 	var files []string
-	var walk func(string) error
-	walk = func(path string) error {
-		f, err := open(path)
-		if err != nil {
-			return err
-		}
-		defer f.Close()
-		fis, err := f.Readdir(-1)
-		if err != nil {
-			files = append(files, path)
-			return nil
-		}
-		for _, fi := range fis {
-			if fi.Name() == "." || fi.Name() == ".." {
-				continue
-			} else if child := path + "/" + fi.Name(); fi.IsDir() {
-				err = walk(child)
-				if err != nil {
-					return err
-				}
-			} else if strings.HasSuffix(child, ".gob") || strings.HasSuffix(child, ".gob.gz") {
-				files = append(files, child)
+	f, err := open(path)
+	if err != nil {
+		return nil, err
+	}
+	defer f.Close()
+	fis, err := f.Readdir(-1)
+	if err != nil {
+		return []string{path}, nil
+	}
+	for _, fi := range fis {
+		if fi.Name() == "." || fi.Name() == ".." {
+			continue
+		} else if child := path + "/" + fi.Name(); fi.IsDir() {
+			add, err := allFiles(child, re)
+			if err != nil {
+				return nil, err
 			}
+			files = append(files, add...)
+		} else if re == nil || re.MatchString(child) {
+			files = append(files, child)
 		}
-		return nil
 	}
+	sort.Strings(files)
+	return files, nil
+}
+
+var matchGobFile = regexp.MustCompile(`\.gob(\.gz)?$`)
+
+func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string) error {
 	log.Infof("LoadDir: walk dir %s", path)
-	err := walk(path)
+	files, err := allFiles(path, matchGobFile)
 	if err != nil {
 		return err
 	}
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 	var mtx sync.Mutex
-	cgs := []CompactGenome{}
-	cseqs := []CompactSequence{}
-	variantmap := map[tileLibRef]tileVariantID{}
+	allcgs := make([][]CompactGenome, len(files))
+	allcseqs := make([][]CompactSequence, len(files))
+	allvariantmap := map[tileLibRef]tileVariantID{}
 	errs := make(chan error, len(files))
 	log.Infof("LoadDir: read %d files", len(files))
-	throttle := throttle{Max: runtime.GOMAXPROCS(0)/2 + 1}
-	for _, path := range files {
-		path := path
+	for fileno, path := range files {
+		fileno, path := fileno, path
 		go func() {
-			throttle.Acquire()
-			defer throttle.Release()
 			f, err := open(path)
 			if err != nil {
 				errs <- err
@@ -264,7 +290,11 @@ func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string, onLoadGeno
 			}
 			defer f.Close()
 			defer log.Infof("LoadDir: finished reading %s", path)
-			errs <- DecodeLibrary(f, strings.HasSuffix(path, ".gz"), func(ent *LibraryEntry) error {
+
+			var variantmap = map[tileLibRef]tileVariantID{}
+			var cgs []CompactGenome
+			var cseqs []CompactSequence
+			err = DecodeLibrary(f, strings.HasSuffix(path, ".gz"), func(ent *LibraryEntry) error {
 				if ctx.Err() != nil {
 					return ctx.Err()
 				}
@@ -281,19 +311,21 @@ func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string, onLoadGeno
 					}
 					mtx.Unlock()
 				}
-				variantmapadd := map[tileLibRef]tileVariantID{}
 				for _, tv := range ent.TileVariants {
-					variantmapadd[tileLibRef{Tag: tv.Tag, Variant: tv.Variant}] = tilelib.getRef(tv.Tag, tv.Sequence).Variant
+					variantmap[tileLibRef{Tag: tv.Tag, Variant: tv.Variant}] = tilelib.getRef(tv.Tag, tv.Sequence, tv.Ref).Variant
 				}
-				mtx.Lock()
 				cgs = append(cgs, ent.CompactGenomes...)
 				cseqs = append(cseqs, ent.CompactSequences...)
-				for k, v := range variantmapadd {
-					variantmap[k] = v
-				}
-				mtx.Unlock()
 				return nil
 			})
+			allcgs[fileno] = cgs
+			allcseqs[fileno] = cseqs
+			mtx.Lock()
+			defer mtx.Unlock()
+			for k, v := range variantmap {
+				allvariantmap[k] = v
+			}
+			errs <- err
 		}()
 	}
 	for range files {
@@ -302,22 +334,34 @@ func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string, onLoadGeno
 			return err
 		}
 	}
+
 	log.Info("LoadDir: loadCompactGenomes")
-	err = tilelib.loadCompactGenomes(cgs, variantmap, onLoadGenome)
+	var flatcgs []CompactGenome
+	for _, cgs := range allcgs {
+		flatcgs = append(flatcgs, cgs...)
+	}
+	err = tilelib.loadCompactGenomes(flatcgs, allvariantmap)
 	if err != nil {
 		return err
 	}
+
 	log.Info("LoadDir: loadCompactSequences")
-	err = tilelib.loadCompactSequences(cseqs, variantmap)
+	var flatcseqs []CompactSequence
+	for _, cseqs := range allcseqs {
+		flatcseqs = append(flatcseqs, cseqs...)
+	}
+	err = tilelib.loadCompactSequences(flatcseqs, allvariantmap)
 	if err != nil {
 		return err
 	}
+
 	log.Info("LoadDir done")
 	return nil
 }
 
 func (tilelib *tileLibrary) WriteDir(dir string) error {
-	nfiles := 128
+	ntilefiles := 128
+	nfiles := ntilefiles + len(tilelib.refseqs)
 	files := make([]*os.File, nfiles)
 	for i := range files {
 		f, err := os.OpenFile(fmt.Sprintf("%s/library.%04d.gob.gz", dir, i), os.O_CREATE|os.O_WRONLY, 0666)
@@ -347,6 +391,12 @@ func (tilelib *tileLibrary) WriteDir(dir string) error {
 	}
 	sort.Strings(cgnames)
 
+	refnames := make([]string, 0, len(tilelib.refseqs))
+	for name := range tilelib.refseqs {
+		refnames = append(refnames, name)
+	}
+	sort.Strings(refnames)
+
 	log.Infof("WriteDir: writing %d files", nfiles)
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
@@ -359,21 +409,17 @@ func (tilelib *tileLibrary) WriteDir(dir string) error {
 				errs <- err
 				return
 			}
-			if start == 0 {
-				// For now, just write all the refs to
-				// the first file
-				for name, tseqs := range tilelib.refseqs {
-					err := encoders[start].Encode(LibraryEntry{CompactSequences: []CompactSequence{{
-						Name:          name,
-						TileSequences: tseqs,
-					}}})
-					if err != nil {
-						errs <- err
-						return
-					}
-				}
+			if refidx := start - ntilefiles; refidx >= 0 {
+				// write each ref to its own file
+				// (they seem to load very slowly)
+				name := refnames[refidx]
+				errs <- encoders[start].Encode(LibraryEntry{CompactSequences: []CompactSequence{{
+					Name:          name,
+					TileSequences: tilelib.refseqs[name],
+				}}})
+				return
 			}
-			for i := start; i < len(cgnames); i += nfiles {
+			for i := start; i < len(cgnames); i += ntilefiles {
 				err := encoders[start].Encode(LibraryEntry{CompactGenomes: []CompactGenome{{
 					Name:     cgnames[i],
 					Variants: tilelib.compactGenomes[cgnames[i]],
@@ -384,14 +430,14 @@ func (tilelib *tileLibrary) WriteDir(dir string) error {
 				}
 			}
 			tvs := []TileVariant{}
-			for tag := start; tag < len(tilelib.variant) && ctx.Err() == nil; tag += nfiles {
+			for tag := start; tag < len(tilelib.variant) && ctx.Err() == nil; tag += ntilefiles {
 				tvs = tvs[:0]
 				for idx, hash := range tilelib.variant[tag] {
 					tvs = append(tvs, TileVariant{
 						Tag:      tagID(tag),
 						Variant:  tileVariantID(idx + 1),
 						Blake2b:  hash,
-						Sequence: tilelib.seq[hash],
+						Sequence: tilelib.hashSequence(hash),
 					})
 				}
 				err := encoders[start].Encode(LibraryEntry{TileVariants: tvs})
@@ -431,9 +477,7 @@ func (tilelib *tileLibrary) WriteDir(dir string) error {
 // Load library data from rdr. Tile variants might be renumbered in
 // the process; in that case, genomes variants will be renumbered to
 // match.
-//
-// If onLoadGenome is non-nil, call it on each CompactGenome entry.
-func (tilelib *tileLibrary) LoadGob(ctx context.Context, rdr io.Reader, gz bool, onLoadGenome func(CompactGenome)) error {
+func (tilelib *tileLibrary) LoadGob(ctx context.Context, rdr io.Reader, gz bool) error {
 	cgs := []CompactGenome{}
 	cseqs := []CompactSequence{}
 	variantmap := map[tileLibRef]tileVariantID{}
@@ -457,7 +501,7 @@ func (tilelib *tileLibrary) LoadGob(ctx context.Context, rdr io.Reader, gz bool,
 	if ctx.Err() != nil {
 		return ctx.Err()
 	}
-	err = tilelib.loadCompactGenomes(cgs, variantmap, onLoadGenome)
+	err = tilelib.loadCompactGenomes(cgs, variantmap)
 	if err != nil {
 		return err
 	}
@@ -507,7 +551,7 @@ type importStats struct {
 	DroppedOutOfOrderTiles int
 }
 
-func (tilelib *tileLibrary) TileFasta(filelabel string, rdr io.Reader, matchChromosome *regexp.Regexp) (tileSeq, []importStats, error) {
+func (tilelib *tileLibrary) TileFasta(filelabel string, rdr io.Reader, matchChromosome *regexp.Regexp, isRef bool) (tileSeq, []importStats, error) {
 	ret := tileSeq{}
 	type jobT struct {
 		label string
@@ -515,6 +559,7 @@ func (tilelib *tileLibrary) TileFasta(filelabel string, rdr io.Reader, matchChro
 	}
 	todo := make(chan jobT, 1)
 	scanner := bufio.NewScanner(rdr)
+	scanner.Buffer(make([]byte, 256), 1<<29) // 512 MiB, in case fasta does not have line breaks
 	go func() {
 		defer close(todo)
 		var fasta []byte
@@ -591,7 +636,7 @@ func (tilelib *tileLibrary) TileFasta(filelabel string, rdr io.Reader, matchChro
 				} else {
 					endpos = found[i+1].pos + taglen
 				}
-				path[i] = tilelib.getRef(f.tagid, job.fasta[startpos:endpos])
+				path[i] = tilelib.getRef(f.tagid, job.fasta[startpos:endpos], isRef)
 				if countBases(job.fasta[startpos:endpos]) != endpos-startpos {
 					atomic.AddInt64(&lowquality, 1)
 				}
@@ -628,7 +673,7 @@ func (tilelib *tileLibrary) Len() int64 {
 
 // Return a tileLibRef for a tile with the given tag and sequence,
 // adding the sequence to the library if needed.
-func (tilelib *tileLibrary) getRef(tag tagID, seq []byte) tileLibRef {
+func (tilelib *tileLibrary) getRef(tag tagID, seq []byte, usedByRef bool) tileLibRef {
 	dropSeq := false
 	if !tilelib.retainNoCalls {
 		for _, b := range seq {
@@ -713,37 +758,66 @@ func (tilelib *tileLibrary) getRef(tag tagID, seq []byte) tileLibRef {
 	vlock.Unlock()
 
 	if tilelib.retainTileSequences && !dropSeq {
-		tilelib.mtx.Lock()
-		if tilelib.seq == nil {
-			tilelib.seq = map[[blake2b.Size256]byte][]byte{}
+		seqCopy := append([]byte(nil), seq...)
+		if tilelib.seq2 == nil {
+			tilelib.mtx.Lock()
+			if tilelib.seq2 == nil {
+				tilelib.seq2lock = map[[2]byte]sync.Locker{}
+				m := map[[2]byte]map[[blake2b.Size256]byte][]byte{}
+				var k [2]byte
+				for i := 0; i < 256; i++ {
+					k[0] = byte(i)
+					for j := 0; j < 256; j++ {
+						k[1] = byte(j)
+						m[k] = map[[blake2b.Size256]byte][]byte{}
+						tilelib.seq2lock[k] = &sync.Mutex{}
+					}
+				}
+				tilelib.seq2 = m
+			}
+			tilelib.mtx.Unlock()
 		}
-		tilelib.seq[seqhash] = append([]byte(nil), seq...)
-		tilelib.mtx.Unlock()
+		var k [2]byte
+		copy(k[:], seqhash[:])
+		locker := tilelib.seq2lock[k]
+		locker.Lock()
+		tilelib.seq2[k][seqhash] = seqCopy
+		locker.Unlock()
 	}
 
+	saveSeq := seq
+	if dropSeq {
+		// Save the hash, but not the sequence
+		saveSeq = nil
+	}
 	if tilelib.encoder != nil {
-		saveSeq := seq
-		if dropSeq {
-			// Save the hash, but not the sequence
-			saveSeq = nil
-		}
 		tilelib.encoder.Encode(LibraryEntry{
 			TileVariants: []TileVariant{{
 				Tag:      tag,
+				Ref:      usedByRef,
 				Variant:  variant,
 				Blake2b:  seqhash,
 				Sequence: saveSeq,
 			}},
 		})
 	}
+	if tilelib.onAddTileVariant != nil {
+		tilelib.onAddTileVariant(tileLibRef{tag, variant}, seqhash, saveSeq)
+	}
 	return tileLibRef{Tag: tag, Variant: variant}
 }
 
+func (tilelib *tileLibrary) hashSequence(hash [blake2b.Size256]byte) []byte {
+	var partition [2]byte
+	copy(partition[:], hash[:])
+	return tilelib.seq2[partition][hash]
+}
+
 func (tilelib *tileLibrary) TileVariantSequence(libref tileLibRef) []byte {
 	if libref.Variant == 0 || len(tilelib.variant) <= int(libref.Tag) || len(tilelib.variant[libref.Tag]) < int(libref.Variant) {
 		return nil
 	}
-	return tilelib.seq[tilelib.variant[libref.Tag][libref.Variant-1]]
+	return tilelib.hashSequence(tilelib.variant[libref.Tag][libref.Variant-1])
 }
 
 // Tidy deletes unreferenced tile variants and renumbers variants so