merge.go

   1 package main
   2
   3 import (
   4         "bufio"
   5         "context"
   6         "encoding/gob"
   7         "errors"
   8         "flag"
   9         "fmt"
  10         "io"
  11         "io/ioutil"
  12         "net/http"
  13         _ "net/http/pprof"
  14         "os"
  15         "sync"
  16
  17         "git.arvados.org/arvados.git/sdk/go/arvados"
  18         log "github.com/sirupsen/logrus"
  19 )
  20
  21 type merger struct {
  22         stdin   io.Reader
  23         inputs  []string
  24         output  io.WriteCloser
  25         tagSet  [][]byte
  26         tilelib *tileLibrary
  27         mapped  map[string]map[tileLibRef]tileVariantID
  28         mtxTags sync.Mutex
  29         errs    chan error
  30 }
  31
  32 func (cmd *merger) 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         pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
  42         runlocal := flags.Bool("local", false, "run on local host (default: run in an arvados container)")
  43         projectUUID := flags.String("project", "", "project `UUID` for output data")
  44         priority := flags.Int("priority", 500, "container request priority")
  45         outputFilename := flags.String("o", "-", "output `file`")
  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         }
  53         cmd.stdin = stdin
  54         cmd.inputs = flags.Args()
  55
  56         if *pprof != "" {
  57                 go func() {
  58                         log.Println(http.ListenAndServe(*pprof, nil))
  59                 }()
  60         }
  61
  62         if !*runlocal {
  63                 if *outputFilename != "-" {
  64                         err = errors.New("cannot specify output file in container mode: not implemented")
  65                         return 1
  66                 }
  67                 runner := arvadosContainerRunner{
  68                         Name:        "lightning filter",
  69                         Client:      arvados.NewClientFromEnv(),
  70                         ProjectUUID: *projectUUID,
  71                         RAM:         64000000000,
  72                         VCPUs:       2,
  73                         Priority:    *priority,
  74                 }
  75                 for i := range cmd.inputs {
  76                         err = runner.TranslatePaths(&cmd.inputs[i])
  77                         if err != nil {
  78                                 return 1
  79                         }
  80                 }
  81                 runner.Args = append([]string{"merge", "-local=true",
  82                         "-o", "/mnt/output/library.gob",
  83                 }, cmd.inputs...)
  84                 var output string
  85                 output, err = runner.Run()
  86                 if err != nil {
  87                         return 1
  88                 }
  89                 fmt.Fprintln(stdout, output+"/library.gob")
  90                 return 0
  91         }
  92
  93         if *outputFilename == "-" {
  94                 cmd.output = nopCloser{stdout}
  95         } else {
  96                 cmd.output, err = os.OpenFile(*outputFilename, os.O_CREATE|os.O_WRONLY, 0777)
  97                 if err != nil {
  98                         return 1
  99                 }
 100                 defer cmd.output.Close()
 101         }
 102
 103         err = cmd.doMerge()
 104         if err != nil {
 105                 return 1
 106         }
 107         err = cmd.output.Close()
 108         if err != nil {
 109                 return 1
 110         }
 111         return 0
 112 }
 113
 114 func (cmd *merger) setError(err error) {
 115         select {
 116         case cmd.errs <- err:
 117         default:
 118         }
 119 }
 120
 121 func (cmd *merger) doMerge() error {
 122         w := bufio.NewWriter(cmd.output)
 123         encoder := gob.NewEncoder(w)
 124
 125         ctx, cancel := context.WithCancel(context.Background())
 126         defer cancel()
 127
 128         cmd.errs = make(chan error, 1)
 129         cmd.tilelib = &tileLibrary{
 130                 encoder:        encoder,
 131                 includeNoCalls: true,
 132         }
 133
 134         cmd.mapped = map[string]map[tileLibRef]tileVariantID{}
 135         for _, input := range cmd.inputs {
 136                 cmd.mapped[input] = map[tileLibRef]tileVariantID{}
 137         }
 138
 139         var wg sync.WaitGroup
 140         for _, input := range cmd.inputs {
 141                 var infile io.ReadCloser
 142                 if input == "-" {
 143                         infile = ioutil.NopCloser(cmd.stdin)
 144                 } else {
 145                         var err error
 146                         infile, err = os.Open(input)
 147                         if err != nil {
 148                                 return err
 149                         }
 150                         defer infile.Close()
 151                 }
 152                 wg.Add(1)
 153                 go func(input string) {
 154                         defer wg.Done()
 155                         log.Printf("%s: reading", input)
 156                         err := cmd.tilelib.LoadGob(ctx, infile, nil)
 157                         if err != nil {
 158                                 cmd.setError(fmt.Errorf("%s: load failed: %w", input, err))
 159                                 cancel()
 160                                 return
 161                         }
 162                         err = infile.Close()
 163                         if err != nil {
 164                                 cmd.setError(fmt.Errorf("%s: error closing input file: %w", input, err))
 165                                 cancel()
 166                                 return
 167                         }
 168                         log.Printf("%s: done", input)
 169                 }(input)
 170         }
 171         wg.Wait()
 172         go close(cmd.errs)
 173         if err := <-cmd.errs; err != nil {
 174                 return err
 175         }
 176         log.Print("flushing")
 177         err := w.Flush()
 178         if err != nil {
 179                 return err
 180         }
 181         return nil
 182 }