import.go

   1 package main
   2
   3 import (
   4         "bufio"
   5         "compress/gzip"
   6         "encoding/gob"
   7         "errors"
   8         "flag"
   9         "fmt"
  10         "io"
  11         "log"
  12         "net/http"
  13         _ "net/http/pprof"
  14         "os"
  15         "os/exec"
  16         "path/filepath"
  17         "regexp"
  18         "runtime"
  19         "sort"
  20         "strings"
  21         "sync"
  22         "sync/atomic"
  23         "time"
  24 )
  25
  26 type importer struct {
  27         tagLibraryFile string
  28         refFile        string
  29         encoder        *gob.Encoder
  30 }
  31
  32 func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
  33         var err error
  34         defer func() {
  35                 if err != nil {
  36                         fmt.Fprintf(stderr, "%s\n", err)
  37                 }
  38         }()
  39         flags := flag.NewFlagSet("", flag.ContinueOnError)
  40         flags.SetOutput(stderr)
  41         flags.StringVar(&cmd.tagLibraryFile, "tag-library", "", "tag library fasta `file`")
  42         flags.StringVar(&cmd.refFile, "ref", "", "reference fasta `file`")
  43         pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
  44         err = flags.Parse(args)
  45         if err == flag.ErrHelp {
  46                 err = nil
  47                 return 0
  48         } else if err != nil {
  49                 return 2
  50         } else if cmd.tagLibraryFile == "" {
  51                 fmt.Fprintln(os.Stderr, "cannot import without -tag-library argument")
  52                 return 2
  53         } else if flags.NArg() == 0 {
  54                 flags.Usage()
  55                 return 2
  56         }
  57
  58         if *pprof != "" {
  59                 go func() {
  60                         log.Println(http.ListenAndServe(*pprof, nil))
  61                 }()
  62         }
  63
  64         infiles, err := listInputFiles(flags.Args())
  65         if err != nil {
  66                 return 1
  67         }
  68
  69         tilelib, err := cmd.loadTileLibrary()
  70         if err != nil {
  71                 return 1
  72         }
  73         go func() {
  74                 for range time.Tick(10 * time.Second) {
  75                         log.Printf("tilelib.Len() == %d", tilelib.Len())
  76                 }
  77         }()
  78         w := bufio.NewWriter(stdout)
  79         cmd.encoder = gob.NewEncoder(w)
  80         err = cmd.tileInputs(tilelib, infiles)
  81         if err != nil {
  82                 return 1
  83         }
  84         err = w.Flush()
  85         if err != nil {
  86                 return 1
  87         }
  88         return 0
  89 }
  90
  91 func (cmd *importer) tileFasta(tilelib *tileLibrary, infile string) (tileSeq, error) {
  92         var input io.ReadCloser
  93         input, err := os.Open(infile)
  94         if err != nil {
  95                 return nil, err
  96         }
  97         defer input.Close()
  98         if strings.HasSuffix(infile, ".gz") {
  99                 input, err = gzip.NewReader(input)
 100                 if err != nil {
 101                         return nil, err
 102                 }
 103                 defer input.Close()
 104         }
 105         return tilelib.TileFasta(infile, input)
 106 }
 107
 108 func (cmd *importer) loadTileLibrary() (*tileLibrary, error) {
 109         log.Printf("tag library %s load starting", cmd.tagLibraryFile)
 110         f, err := os.Open(cmd.tagLibraryFile)
 111         if err != nil {
 112                 return nil, err
 113         }
 114         defer f.Close()
 115         var rdr io.ReadCloser = f
 116         if strings.HasSuffix(cmd.tagLibraryFile, ".gz") {
 117                 rdr, err = gzip.NewReader(f)
 118                 if err != nil {
 119                         return nil, fmt.Errorf("%s: gzip: %s", cmd.tagLibraryFile, err)
 120                 }
 121                 defer rdr.Close()
 122         }
 123         var taglib tagLibrary
 124         err = taglib.Load(rdr)
 125         if err != nil {
 126                 return nil, err
 127         }
 128         if taglib.Len() < 1 {
 129                 return nil, fmt.Errorf("cannot tile: tag library is empty")
 130         }
 131         log.Printf("tag library %s load done", cmd.tagLibraryFile)
 132         return &tileLibrary{taglib: &taglib}, nil
 133 }
 134
 135 func listInputFiles(paths []string) (files []string, err error) {
 136         for _, path := range paths {
 137                 if fi, err := os.Stat(path); err != nil {
 138                         return nil, fmt.Errorf("%s: stat failed: %s", path, err)
 139                 } else if !fi.IsDir() {
 140                         if !strings.HasSuffix(path, ".2.fasta") || strings.HasSuffix(path, ".2.fasta.gz") {
 141                                 files = append(files, path)
 142                         }
 143                         continue
 144                 }
 145                 d, err := os.Open(path)
 146                 if err != nil {
 147                         return nil, fmt.Errorf("%s: open failed: %s", path, err)
 148                 }
 149                 defer d.Close()
 150                 names, err := d.Readdirnames(0)
 151                 if err != nil {
 152                         return nil, fmt.Errorf("%s: readdir failed: %s", path, err)
 153                 }
 154                 sort.Strings(names)
 155                 for _, name := range names {
 156                         if strings.HasSuffix(name, ".vcf") || strings.HasSuffix(name, ".vcf.gz") {
 157                                 files = append(files, filepath.Join(path, name))
 158                         } else if strings.HasSuffix(name, ".1.fasta") || strings.HasSuffix(name, ".1.fasta.gz") {
 159                                 files = append(files, filepath.Join(path, name))
 160                         }
 161                 }
 162                 d.Close()
 163         }
 164         for _, file := range files {
 165                 if strings.HasSuffix(file, ".1.fasta") || strings.HasSuffix(file, ".1.fasta.gz") {
 166                         continue
 167                 } else if _, err := os.Stat(file + ".csi"); err == nil {
 168                         continue
 169                 } else if _, err = os.Stat(file + ".tbi"); err == nil {
 170                         continue
 171                 } else {
 172                         return nil, fmt.Errorf("%s: cannot read without .tbi or .csi index file", file)
 173                 }
 174         }
 175         return
 176 }
 177
 178 func (cmd *importer) tileInputs(tilelib *tileLibrary, infiles []string) error {
 179         starttime := time.Now()
 180         errs := make(chan error, 1)
 181         todo := make(chan func() error, len(infiles)*2)
 182         var encodeJobs sync.WaitGroup
 183         for _, infile := range infiles {
 184                 infile := infile
 185                 var phases sync.WaitGroup
 186                 phases.Add(2)
 187                 variants := make([][]tileVariantID, 2)
 188                 if strings.HasSuffix(infile, ".1.fasta") || strings.HasSuffix(infile, ".1.fasta.gz") {
 189                         todo <- func() error {
 190                                 defer phases.Done()
 191                                 log.Printf("%s starting", infile)
 192                                 defer log.Printf("%s done", infile)
 193                                 tseqs, err := cmd.tileFasta(tilelib, infile)
 194                                 variants[0] = tseqs.Variants()
 195                                 return err
 196                         }
 197                         infile2 := regexp.MustCompile(`\.1\.fasta(\.gz)?$`).ReplaceAllString(infile, `.2.fasta$1`)
 198                         todo <- func() error {
 199                                 defer phases.Done()
 200                                 log.Printf("%s starting", infile2)
 201                                 defer log.Printf("%s done", infile2)
 202                                 tseqs, err := cmd.tileFasta(tilelib, infile2)
 203                                 variants[1] = tseqs.Variants()
 204                                 return err
 205                         }
 206                 } else {
 207                         for phase := 0; phase < 2; phase++ {
 208                                 phase := phase
 209                                 todo <- func() error {
 210                                         defer phases.Done()
 211                                         log.Printf("%s phase %d starting", infile, phase+1)
 212                                         defer log.Printf("%s phase %d done", infile, phase+1)
 213                                         tseqs, err := cmd.tileGVCF(tilelib, infile, phase)
 214                                         variants[phase] = tseqs.Variants()
 215                                         return err
 216                                 }
 217                         }
 218                 }
 219                 encodeJobs.Add(1)
 220                 go func() {
 221                         defer encodeJobs.Done()
 222                         phases.Wait()
 223                         if len(errs) > 0 {
 224                                 return
 225                         }
 226                         ntags := len(variants[0])
 227                         if ntags < len(variants[1]) {
 228                                 ntags = len(variants[1])
 229                         }
 230                         flat := make([]tileVariantID, ntags*2)
 231                         for i := 0; i < ntags; i++ {
 232                                 flat[i*2] = variants[0][i]
 233                                 flat[i*2+1] = variants[1][i]
 234                         }
 235                         err := cmd.encoder.Encode(LibraryEntry{
 236                                 CompactGenomes: []CompactGenome{{Name: infile, Variants: flat}},
 237                         })
 238                         if err != nil {
 239                                 select {
 240                                 case errs <- err:
 241                                 default:
 242                                 }
 243                         }
 244                 }()
 245         }
 246         go close(todo)
 247         var tileJobs sync.WaitGroup
 248         running := int64(runtime.NumCPU())
 249         for i := 0; i < runtime.NumCPU(); i++ {
 250                 tileJobs.Add(1)
 251                 go func() {
 252                         defer tileJobs.Done()
 253                         defer atomic.AddInt64(&running, -1)
 254                         for fn := range todo {
 255                                 if len(errs) > 0 {
 256                                         return
 257                                 }
 258                                 err := fn()
 259                                 if err != nil {
 260                                         select {
 261                                         case errs <- err:
 262                                         default:
 263                                         }
 264                                 }
 265                                 remain := len(todo) + int(atomic.LoadInt64(&running)) - 1
 266                                 ttl := time.Now().Sub(starttime) * time.Duration(remain) / time.Duration(cap(todo)-remain)
 267                                 eta := time.Now().Add(ttl)
 268                                 log.Printf("progress %d/%d, eta %v (%v)", cap(todo)-remain, cap(todo), eta, ttl)
 269                         }
 270                 }()
 271         }
 272         tileJobs.Wait()
 273         encodeJobs.Wait()
 274         go close(errs)
 275         return <-errs
 276 }
 277
 278 func (cmd *importer) tileGVCF(tilelib *tileLibrary, infile string, phase int) (tileseq tileSeq, err error) {
 279         if cmd.refFile == "" {
 280                 err = errors.New("cannot import vcf: reference data (-ref) not specified")
 281                 return
 282         }
 283         args := []string{"bcftools", "consensus", "--fasta-ref", cmd.refFile, "-H", fmt.Sprint(phase + 1), infile}
 284         indexsuffix := ".tbi"
 285         if _, err := os.Stat(infile + ".csi"); err == nil {
 286                 indexsuffix = ".csi"
 287         }
 288         if out, err := exec.Command("docker", "image", "ls", "-q", "lightning-runtime").Output(); err == nil && len(out) > 0 {
 289                 args = append([]string{
 290                         "docker", "run", "--rm",
 291                         "--log-driver=none",
 292                         "--volume=" + infile + ":" + infile + ":ro",
 293                         "--volume=" + infile + indexsuffix + ":" + infile + indexsuffix + ":ro",
 294                         "--volume=" + cmd.refFile + ":" + cmd.refFile + ":ro",
 295                         "lightning-runtime",
 296                 }, args...)
 297         }
 298         consensus := exec.Command(args[0], args[1:]...)
 299         consensus.Stderr = os.Stderr
 300         stdout, err := consensus.StdoutPipe()
 301         defer stdout.Close()
 302         if err != nil {
 303                 return
 304         }
 305         err = consensus.Start()
 306         if err != nil {
 307                 return
 308         }
 309         defer consensus.Wait()
 310         tileseq, err = tilelib.TileFasta(fmt.Sprintf("%s phase %d", infile, phase+1), stdout)
 311         if err != nil {
 312                 return
 313         }
 314         err = stdout.Close()
 315         if err != nil {
 316                 return
 317         }
 318         err = consensus.Wait()
 319         if err != nil {
 320                 err = fmt.Errorf("%s phase %d: bcftools: %s", infile, phase, err)
 321                 return
 322         }
 323         return
 324 }