stats.go

   1 package main
   2
   3 import (
   4         "bufio"
   5         "encoding/gob"
   6         "encoding/json"
   7         "errors"
   8         "flag"
   9         "fmt"
  10         "io"
  11         "io/ioutil"
  12         "net/http"
  13         _ "net/http/pprof"
  14         "os"
  15
  16         "git.arvados.org/arvados.git/sdk/go/arvados"
  17         log "github.com/sirupsen/logrus"
  18 )
  19
  20 type stats struct{}
  21
  22 func (cmd *stats) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
  23         var err error
  24         defer func() {
  25                 if err != nil {
  26                         fmt.Fprintf(stderr, "%s\n", err)
  27                 }
  28         }()
  29         flags := flag.NewFlagSet("", flag.ContinueOnError)
  30         flags.SetOutput(stderr)
  31         pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
  32         runlocal := flags.Bool("local", false, "run on local host (default: run in an arvados container)")
  33         projectUUID := flags.String("project", "", "project `UUID` for output data")
  34         priority := flags.Int("priority", 500, "container request priority")
  35         inputFilename := flags.String("i", "-", "input `file`")
  36         outputFilename := flags.String("o", "-", "output `file`")
  37         err = flags.Parse(args)
  38         if err == flag.ErrHelp {
  39                 err = nil
  40                 return 0
  41         } else if err != nil {
  42                 return 2
  43         }
  44
  45         if *pprof != "" {
  46                 go func() {
  47                         log.Println(http.ListenAndServe(*pprof, nil))
  48                 }()
  49         }
  50
  51         if !*runlocal {
  52                 if *outputFilename != "-" {
  53                         err = errors.New("cannot specify output file in container mode: not implemented")
  54                         return 1
  55                 }
  56                 runner := arvadosContainerRunner{
  57                         Name:        "lightning stats",
  58                         Client:      arvados.NewClientFromEnv(),
  59                         ProjectUUID: *projectUUID,
  60                         RAM:         16000000000,
  61                         VCPUs:       1,
  62                         Priority:    *priority,
  63                 }
  64                 err = runner.TranslatePaths(inputFilename)
  65                 if err != nil {
  66                         return 1
  67                 }
  68                 runner.Args = []string{"stats", "-local=true", "-i", *inputFilename, "-o", "/mnt/output/stats.json"}
  69                 var output string
  70                 output, err = runner.Run()
  71                 if err != nil {
  72                         return 1
  73                 }
  74                 fmt.Fprintln(stdout, output+"/stats.json")
  75                 return 0
  76         }
  77
  78         var input io.ReadCloser
  79         if *inputFilename == "-" {
  80                 input = ioutil.NopCloser(stdin)
  81         } else {
  82                 input, err = os.Open(*inputFilename)
  83                 if err != nil {
  84                         return 1
  85                 }
  86                 defer input.Close()
  87         }
  88
  89         var output io.WriteCloser
  90         if *outputFilename == "-" {
  91                 output = nopCloser{stdout}
  92         } else {
  93                 output, err = os.OpenFile(*outputFilename, os.O_CREATE|os.O_WRONLY, 0777)
  94                 if err != nil {
  95                         return 1
  96                 }
  97                 defer output.Close()
  98         }
  99
 100         bufw := bufio.NewWriter(output)
 101         cmd.readLibrary(input, bufw)
 102         err = bufw.Flush()
 103         if err != nil {
 104                 return 1
 105         }
 106         err = output.Close()
 107         if err != nil {
 108                 return 1
 109         }
 110         return 0
 111 }
 112
 113 func (cmd *stats) readLibrary(input io.Reader, output io.Writer) error {
 114         var ret struct {
 115                 Genomes          int
 116                 Tags             int
 117                 TileVariants     int
 118                 VariantsBySize   []int
 119                 NCVariantsBySize []int
 120         }
 121
 122         dec := gob.NewDecoder(input)
 123         for {
 124                 var ent LibraryEntry
 125                 err := dec.Decode(&ent)
 126                 if err == io.EOF {
 127                         break
 128                 } else if err != nil {
 129                         return err
 130                 }
 131                 ret.Genomes += len(ent.CompactGenomes)
 132                 ret.Tags += len(ent.TagSet)
 133                 ret.TileVariants += len(ent.TileVariants)
 134                 for _, tv := range ent.TileVariants {
 135                         if need := 1 + len(tv.Sequence) - len(ret.VariantsBySize); need > 0 {
 136                                 ret.VariantsBySize = append(ret.VariantsBySize, make([]int, need)...)
 137                                 ret.NCVariantsBySize = append(ret.NCVariantsBySize, make([]int, need)...)
 138                         }
 139
 140                         hasNoCalls := false
 141                         for _, b := range tv.Sequence {
 142                                 if b != 'a' && b != 'c' && b != 'g' && b != 't' {
 143                                         hasNoCalls = true
 144                                 }
 145                         }
 146
 147                         if hasNoCalls {
 148                                 ret.NCVariantsBySize[len(tv.Sequence)]++
 149                         } else {
 150                                 ret.VariantsBySize[len(tv.Sequence)]++
 151                         }
 152                 }
 153         }
 154         return json.NewEncoder(output).Encode(ret)
 155 }