Fix array index out of bounds.

[lightning.git] / arvados.go

diff --git a/arvados.go b/arvados.go
index f2819d01ef2c8ced3b4e30f0076bfa42aa3fa49c..c952a9870c77cab73c35237440555fb2b551ce23 100644 (file)
--- a/arvados.go
+++ b/arvados.go
@@ -1,6 +1,11 @@
-package main
+// Copyright (C) The Lightning Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
+package lightning
 
 import (
+       "bufio"
        "context"
        "encoding/json"
        "errors"
@@ -18,6 +23,7 @@ import (
        "git.arvados.org/arvados.git/sdk/go/arvados"
        "git.arvados.org/arvados.git/sdk/go/arvadosclient"
        "git.arvados.org/arvados.git/sdk/go/keepclient"
+       "github.com/klauspost/pgzip"
        log "github.com/sirupsen/logrus"
        "golang.org/x/crypto/blake2b"
        "golang.org/x/net/websocket"
@@ -261,6 +267,13 @@ func (runner *arvadosContainerRunner) RunContext(ctx context.Context) (string, e
                        "runtime_constraints": rc,
                        "priority":            runner.Priority,
                        "state":               arvados.ContainerRequestStateCommitted,
+                       "scheduling_parameters": arvados.SchedulingParameters{
+                               Preemptible: true,
+                               Partitions:  []string{},
+                       },
+                       "environment": map[string]string{
+                               "GOMAXPROCS": fmt.Sprintf("%d", rc.VCPUs),
+                       },
                },
        })
        if err != nil {
@@ -275,6 +288,7 @@ func (runner *arvadosContainerRunner) RunContext(ctx context.Context) (string, e
        subscribedUUID := ""
        defer func() {
                if subscribedUUID != "" {
+                       log.Printf("unsubscribe container UUID: %s", subscribedUUID)
                        client.Unsubscribe(logch, subscribedUUID)
                }
        }()
@@ -298,8 +312,10 @@ func (runner *arvadosContainerRunner) RunContext(ctx context.Context) (string, e
                        fmt.Fprint(os.Stderr, neednewline)
                        neednewline = ""
                        if subscribedUUID != "" {
+                               log.Printf("unsubscribe container UUID: %s", subscribedUUID)
                                client.Unsubscribe(logch, subscribedUUID)
                        }
+                       log.Printf("subscribe container UUID: %s", cr.ContainerUUID)
                        client.Subscribe(logch, cr.ContainerUUID)
                        subscribedUUID = cr.ContainerUUID
                }
@@ -463,9 +479,51 @@ func (runner *arvadosContainerRunner) makeCommandCollection() (string, error) {
        return coll.UUID, nil
 }
 
-var arvadosClientFromEnv = arvados.NewClientFromEnv()
+// zopen returns a reader for the given file, using the arvados API
+// instead of arv-mount/fuse where applicable, and transparently
+// decompressing the input if fnm ends with ".gz".
+func zopen(fnm string) (io.ReadCloser, error) {
+       f, err := open(fnm)
+       if err != nil || !strings.HasSuffix(fnm, ".gz") {
+               return f, err
+       }
+       rdr, err := pgzip.NewReader(bufio.NewReaderSize(f, 4*1024*1024))
+       if err != nil {
+               f.Close()
+               return nil, err
+       }
+       return gzipr{rdr, f}, nil
+}
+
+// gzipr wraps a ReadCloser and a Closer, presenting a single Close()
+// method that closes both wrapped objects.
+type gzipr struct {
+       io.ReadCloser
+       io.Closer
+}
+
+func (gr gzipr) Close() error {
+       e1 := gr.ReadCloser.Close()
+       e2 := gr.Closer.Close()
+       if e1 != nil {
+               return e1
+       }
+       return e2
+}
+
+var (
+       arvadosClientFromEnv = arvados.NewClientFromEnv()
+       keepClient           *keepclient.KeepClient
+       siteFS               arvados.CustomFileSystem
+       siteFSMtx            sync.Mutex
+)
 
-func open(fnm string) (io.ReadCloser, error) {
+type file interface {
+       io.ReadCloser
+       Readdir(n int) ([]os.FileInfo, error)
+}
+
+func open(fnm string) (file, error) {
        if os.Getenv("ARVADOS_API_HOST") == "" {
                return os.Open(fnm)
        }
@@ -474,31 +532,29 @@ func open(fnm string) (io.ReadCloser, error) {
                return os.Open(fnm)
        }
        uuid := m[2]
-       mnt := "/mnt/" + uuid + "/"
-       if !strings.HasPrefix(fnm, mnt) {
+       mnt := "/mnt/" + uuid
+       if fnm != mnt && !strings.HasPrefix(fnm, mnt+"/") {
                return os.Open(fnm)
        }
 
-       log.Infof("reading %q from %s using Arvados client library", fnm[len(mnt):], uuid)
-       ac, err := arvadosclient.New(arvadosClientFromEnv)
-       if err != nil {
-               return nil, err
-       }
-       ac.Client = arvados.DefaultSecureClient
-       kc := keepclient.New(ac)
-       // Don't use keepclient's default short timeouts.
-       kc.HTTPClient = arvados.DefaultSecureClient
-       // Don't cache more than one block for this file.
-       kc.BlockCache = &keepclient.BlockCache{MaxBlocks: 1}
-
-       var coll arvados.Collection
-       err = arvadosClientFromEnv.RequestAndDecode(&coll, "GET", "arvados/v1/collections/"+uuid, nil, arvados.GetOptions{Select: []string{"uuid", "manifest_text"}})
-       if err != nil {
-               return nil, err
-       }
-       fs, err := coll.FileSystem(arvadosClientFromEnv, kc)
-       if err != nil {
-               return nil, err
-       }
-       return fs.Open(fnm[len(mnt):])
+       siteFSMtx.Lock()
+       defer siteFSMtx.Unlock()
+       if siteFS == nil {
+               log.Info("setting up Arvados client")
+               ac, err := arvadosclient.New(arvadosClientFromEnv)
+               if err != nil {
+                       return nil, err
+               }
+               ac.Client = arvados.DefaultSecureClient
+               keepClient = keepclient.New(ac)
+               // Don't use keepclient's default short timeouts.
+               keepClient.HTTPClient = arvados.DefaultSecureClient
+               keepClient.BlockCache = &keepclient.BlockCache{MaxBlocks: 4}
+               siteFS = arvadosClientFromEnv.SiteFileSystem(keepClient)
+       } else {
+               keepClient.BlockCache.MaxBlocks++
+       }
+
+       log.Infof("reading %q from %s using Arvados client", fnm[len(mnt):], uuid)
+       return siteFS.Open("by_id/" + uuid + fnm[len(mnt):])
 }