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