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

diff --git a/tilelib.go b/tilelib.go
index 8073eba183..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 (
@@ -67,6 +71,12 @@ type tileLibrary struct {
 	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)
@@ -146,8 +156,15 @@ func (tilelib *tileLibrary) loadCompactGenomes(cgs []CompactGenome, variantmap m
 				// 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,39 +232,43 @@ 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
 	}
@@ -283,7 +312,7 @@ func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string, onLoadGeno
 					mtx.Unlock()
 				}
 				for _, tv := range ent.TileVariants {
-					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
 				}
 				cgs = append(cgs, ent.CompactGenomes...)
 				cseqs = append(cseqs, ent.CompactSequences...)
@@ -311,7 +340,7 @@ func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string, onLoadGeno
 	for _, cgs := range allcgs {
 		flatcgs = append(flatcgs, cgs...)
 	}
-	err = tilelib.loadCompactGenomes(flatcgs, allvariantmap, onLoadGenome)
+	err = tilelib.loadCompactGenomes(flatcgs, allvariantmap)
 	if err != nil {
 		return err
 	}
@@ -331,7 +360,8 @@ func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string, onLoadGeno
 }
 
 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)
@@ -361,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()
@@ -373,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]],
@@ -398,7 +430,7 @@ 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{
@@ -445,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{}
@@ -471,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
 	}
@@ -521,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
@@ -529,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
@@ -605,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)
 				}
@@ -642,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 {
@@ -754,21 +785,25 @@ func (tilelib *tileLibrary) getRef(tag tagID, seq []byte) tileLibRef {
 		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}
 }