dumpgob.go

   1 // Copyright (C) The Lightning Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package lightning
   6
   7 import (
   8         "bufio"
   9         "errors"
  10         "flag"
  11         "fmt"
  12         "io"
  13         "net/http"
  14         _ "net/http/pprof"
  15         "os"
  16         "strings"
  17
  18         "git.arvados.org/arvados.git/sdk/go/arvados"
  19         log "github.com/sirupsen/logrus"
  20 )
  21
  22 type dumpGob struct{}
  23
  24 func (cmd *dumpGob) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
  25         var err error
  26         defer func() {
  27                 if err != nil {
  28                         fmt.Fprintf(stderr, "%s\n", err)
  29                 }
  30         }()
  31         flags := flag.NewFlagSet("", flag.ContinueOnError)
  32         flags.SetOutput(stderr)
  33         pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
  34         runlocal := flags.Bool("local", false, "run on local host (default: run in an arvados container)")
  35         projectUUID := flags.String("project", "", "project `UUID` for output data")
  36         priority := flags.Int("priority", 500, "container request priority")
  37         inputFilename := flags.String("i", "-", "input `file` (library)")
  38         outputFilename := flags.String("o", "-", "output `file`")
  39         err = flags.Parse(args)
  40         if err == flag.ErrHelp {
  41                 err = nil
  42                 return 0
  43         } else if err != nil {
  44                 return 2
  45         } else if flags.NArg() > 0 {
  46                 err = fmt.Errorf("errant command line arguments after parsed flags: %v", flags.Args())
  47                 return 2
  48         }
  49
  50         if *pprof != "" {
  51                 go func() {
  52                         log.Println(http.ListenAndServe(*pprof, nil))
  53                 }()
  54         }
  55
  56         if !*runlocal {
  57                 if *outputFilename != "-" {
  58                         err = errors.New("cannot specify output file in container mode: not implemented")
  59                         return 1
  60                 }
  61                 runner := arvadosContainerRunner{
  62                         Name:        "lightning dumpgob",
  63                         Client:      arvados.NewClientFromEnv(),
  64                         ProjectUUID: *projectUUID,
  65                         RAM:         4000000000,
  66                         VCPUs:       1,
  67                         Priority:    *priority,
  68                 }
  69                 err = runner.TranslatePaths(inputFilename)
  70                 if err != nil {
  71                         return 1
  72                 }
  73                 runner.Args = []string{"dumpgob", "-local=true", fmt.Sprintf("-pprof=%v", *pprof), "-i", *inputFilename, "-o", "/mnt/output/dumpgob.txt"}
  74                 var output string
  75                 output, err = runner.Run()
  76                 if err != nil {
  77                         return 1
  78                 }
  79                 fmt.Fprintln(stdout, output+"/dumpgob.txt")
  80                 return 0
  81         }
  82
  83         input, err := open(*inputFilename)
  84         if err != nil {
  85                 return 1
  86         }
  87         defer input.Close()
  88         output, err := os.OpenFile(*outputFilename, os.O_CREATE|os.O_WRONLY, 0644)
  89         if err != nil {
  90                 return 1
  91         }
  92         defer output.Close()
  93         bufw := bufio.NewWriterSize(output, 8*1024*1024)
  94
  95         var n, nCG, nCS, nTV int
  96         err = DecodeLibrary(input, strings.HasSuffix(*inputFilename, ".gz"), func(ent *LibraryEntry) error {
  97                 if n%1000000 == 0 {
  98                         fmt.Fprintf(stderr, "ent %d\n", n)
  99                 }
 100                 n++
 101                 if len(ent.TagSet) > 0 {
 102                         fmt.Fprintf(bufw, "ent %d: TagSet, len %d, taglen %d\n", n, len(ent.TagSet), len(ent.TagSet[0]))
 103                 }
 104                 for _, cg := range ent.CompactGenomes {
 105                         nCG++
 106                         fmt.Fprintf(bufw, "ent %d: CompactGenome, name %q, len(Variants) %d\n", n, cg.Name, len(cg.Variants))
 107                 }
 108                 for _, cs := range ent.CompactSequences {
 109                         nCS++
 110                         fmt.Fprintf(bufw, "ent %d: CompactSequence, name %q, len(TileSequences) %d\n", n, cs.Name, len(cs.TileSequences))
 111                 }
 112                 for _, tv := range ent.TileVariants {
 113                         nTV++
 114                         fmt.Fprintf(bufw, "ent %d: TileVariant, tag %d, variant %d, hash %x, len(seq) %d\n", n, tv.Tag, tv.Variant, tv.Blake2b, len(tv.Sequence))
 115                 }
 116                 return nil
 117         })
 118         if err != nil {
 119                 return 1
 120         }
 121         fmt.Fprintf(bufw, "total: ents %d, CompactGenomes %d, CompactSequences %d, TileVariants %d\n", n, nCG, nCS, nTV)
 122         err = bufw.Flush()
 123         if err != nil {
 124                 return 1
 125         }
 126         err = output.Close()
 127         if err != nil {
 128                 return 1
 129         }
 130         return 0
 131 }