X-Git-Url: https://git.arvados.org/lightning.git/blobdiff_plain/de10f554682eecf3926b409c70f9685901d11272..ed9183d5e6ce23f1cf6f9a7d65d3e68214393e8a:/tilelib.go diff --git a/tilelib.go b/tilelib.go index 926c0324a2..289f8895e5 100644 --- a/tilelib.go +++ b/tilelib.go @@ -1,4 +1,8 @@ -package main +// Copyright (C) The Lightning Authors. All rights reserved. +// +// SPDX-License-Identifier: AGPL-3.0 + +package lightning import ( "bufio" @@ -7,6 +11,7 @@ import ( "encoding/gob" "fmt" "io" + "os" "regexp" "runtime" "sort" @@ -14,6 +19,7 @@ import ( "sync" "sync/atomic" + "github.com/klauspost/pgzip" log "github.com/sirupsen/logrus" "golang.org/x/crypto/blake2b" ) @@ -60,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 @@ -110,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) @@ -141,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{ @@ -172,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 { @@ -196,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() @@ -205,15 +232,249 @@ func (tilelib *tileLibrary) loadCompactSequences(cseqs []CompactSequence, varian for _, cseq := range cseqs { tilelib.refseqs[cseq.Name] = cseq.TileSequences } + log.Info("loadCompactSequences: done") + return nil +} + +func allGobFiles(path string) ([]string, error) { + var files []string + 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 := allGobFiles(child) + if err != nil { + return nil, err + } + files = append(files, add...) + } else if strings.HasSuffix(child, ".gob") || strings.HasSuffix(child, ".gob.gz") { + files = append(files, child) + } + } + return files, nil +} + +func (tilelib *tileLibrary) LoadDir(ctx context.Context, path string) error { + log.Infof("LoadDir: walk dir %s", path) + files, err := allGobFiles(path) + if err != nil { + return err + } + ctx, cancel := context.WithCancel(ctx) + defer cancel() + var mtx sync.Mutex + 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)) + for fileno, path := range files { + fileno, path := fileno, path + go func() { + f, err := open(path) + if err != nil { + errs <- err + return + } + defer f.Close() + defer log.Infof("LoadDir: finished reading %s", path) + + 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() + } + if len(ent.TagSet) > 0 { + mtx.Lock() + if tilelib.taglib == nil || tilelib.taglib.Len() != len(ent.TagSet) { + // load first set of tags, or + // report mismatch if 2 sets + // have different #tags. + if err := tilelib.loadTagSet(ent.TagSet); err != nil { + mtx.Unlock() + return err + } + } + mtx.Unlock() + } + for _, tv := range ent.TileVariants { + 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...) + 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 { + err := <-errs + if err != nil { + return err + } + } + + log.Info("LoadDir: loadCompactGenomes") + 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") + 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 { + 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) + if err != nil { + return err + } + defer f.Close() + files[i] = f + } + bufws := make([]*bufio.Writer, nfiles) + for i := range bufws { + bufws[i] = bufio.NewWriterSize(files[i], 1<<26) + } + zws := make([]*pgzip.Writer, nfiles) + for i := range zws { + zws[i] = pgzip.NewWriter(bufws[i]) + defer zws[i].Close() + } + encoders := make([]*gob.Encoder, nfiles) + for i := range encoders { + encoders[i] = gob.NewEncoder(zws[i]) + } + + cgnames := make([]string, 0, len(tilelib.compactGenomes)) + for name := range tilelib.compactGenomes { + cgnames = append(cgnames, name) + } + 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() + errs := make(chan error, nfiles) + for start := range files { + start := start + go func() { + err := encoders[start].Encode(LibraryEntry{TagSet: tilelib.taglib.Tags()}) + 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 += ntilefiles { + err := encoders[start].Encode(LibraryEntry{CompactGenomes: []CompactGenome{{ + Name: cgnames[i], + Variants: tilelib.compactGenomes[cgnames[i]], + }}}) + if err != nil { + errs <- err + return + } + } + tvs := []TileVariant{} + 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.hashSequence(hash), + }) + } + err := encoders[start].Encode(LibraryEntry{TileVariants: tvs}) + if err != nil { + errs <- err + return + } + } + errs <- nil + }() + } + for range files { + err := <-errs + if err != nil { + return err + } + } + log.Info("WriteDir: flushing") + for i := range zws { + err := zws[i].Close() + if err != nil { + return err + } + err = bufws[i].Flush() + if err != nil { + return err + } + err = files[i].Close() + if err != nil { + return err + } + } + log.Info("WriteDir: done") return nil } // 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{} @@ -237,7 +498,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 } @@ -248,6 +509,36 @@ func (tilelib *tileLibrary) LoadGob(ctx context.Context, rdr io.Reader, gz bool, return nil } +func (tilelib *tileLibrary) dump(out io.Writer) { + printTV := func(tag int, variant tileVariantID) { + if variant < 1 { + fmt.Fprintf(out, " -") + } else if tag >= len(tilelib.variant) { + fmt.Fprintf(out, " (!tag=%d)", tag) + } else if int(variant) > len(tilelib.variant[tag]) { + fmt.Fprintf(out, " (tag=%d,!variant=%d)", tag, variant) + } else { + fmt.Fprintf(out, " %x", tilelib.variant[tag][variant-1][:8]) + } + } + for refname, refseqs := range tilelib.refseqs { + for seqname, seq := range refseqs { + fmt.Fprintf(out, "ref %s %s", refname, seqname) + for _, libref := range seq { + printTV(int(libref.Tag), libref.Variant) + } + fmt.Fprintf(out, "\n") + } + } + for name, cg := range tilelib.compactGenomes { + fmt.Fprintf(out, "cg %s", name) + for tag, variant := range cg { + printTV(tag/2, variant) + } + fmt.Fprintf(out, "\n") + } +} + type importStats struct { InputFile string InputLabel string @@ -257,7 +548,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 @@ -324,6 +615,7 @@ func (tilelib *tileLibrary) TileFasta(filelabel string, rdr io.Reader, matchChro log.Infof("%s %s getting %d librefs", filelabel, job.label, len(found)) throttle := &throttle{Max: runtime.NumCPU()} path = path[:len(found)] + var lowquality int64 for i, f := range found { i, f := i, f throttle.Acquire() @@ -340,7 +632,10 @@ 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) + } }() } throttle.Wait() @@ -352,7 +647,7 @@ func (tilelib *tileLibrary) TileFasta(filelabel string, rdr io.Reader, matchChro ret[job.label] = pathcopy basesIn := countBases(job.fasta) - log.Infof("%s %s fasta in %d coverage in %d path len %d skipped-out-of-order %d", filelabel, job.label, len(job.fasta), basesIn, len(path), skipped) + log.Infof("%s %s fasta in %d coverage in %d path len %d low-quality %d skipped-out-of-order %d", filelabel, job.label, len(job.fasta), basesIn, len(path), lowquality, skipped) stats = append(stats, importStats{ InputFile: filelabel, InputLabel: job.label, @@ -374,7 +669,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 { @@ -459,37 +754,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 @@ -560,18 +884,30 @@ func (tilelib *tileLibrary) Tidy() { // Apply remap to genomes and reference sequences, so they // refer to the same tile variants using the changed IDs. log.Print("Tidy: apply remap") + var wg sync.WaitGroup for _, cg := range tilelib.compactGenomes { - for idx, variant := range cg { - cg[idx] = remap[tagID(idx/2)][variant] - } + cg := cg + wg.Add(1) + go func() { + defer wg.Done() + for idx, variant := range cg { + cg[idx] = remap[tagID(idx/2)][variant] + } + }() } for _, refcs := range tilelib.refseqs { for _, refseq := range refcs { - for i, tv := range refseq { - refseq[i].Variant = remap[tv.Tag][tv.Variant] - } + refseq := refseq + wg.Add(1) + go func() { + defer wg.Done() + for i, tv := range refseq { + refseq[i].Variant = remap[tv.Tag][tv.Variant] + } + }() } } + wg.Wait() log.Print("Tidy: done") }