X-Git-Url: https://git.arvados.org/lightning.git/blobdiff_plain/a0d7eb079fe583a9c8895a848cd3a01f786a9a4c..8a9c807d4b80f61a8485c0fc07abfe23f9d50202:/arvados.go

diff --git a/arvados.go b/arvados.go
index 89d23c7e31..cc41ba2c9a 100644
--- a/arvados.go
+++ b/arvados.go
@@ -1,19 +1,181 @@
 package main
 
 import (
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io/ioutil"
-	"log"
+	"net/url"
 	"os"
 	"regexp"
+	"strings"
+	"sync"
+	"time"
 
 	"git.arvados.org/arvados.git/sdk/go/arvados"
 	"git.arvados.org/arvados.git/sdk/go/arvadosclient"
 	"git.arvados.org/arvados.git/sdk/go/keepclient"
+	log "github.com/sirupsen/logrus"
 	"golang.org/x/crypto/blake2b"
+	"golang.org/x/net/websocket"
 )
 
+type eventMessage struct {
+	Status     int
+	ObjectUUID string `json:"object_uuid"`
+	EventType  string `json:"event_type"`
+	Properties struct {
+		Text string
+	}
+}
+
+type arvadosClient struct {
+	*arvados.Client
+	notifying map[string]map[chan<- eventMessage]int
+	wantClose chan struct{}
+	wsconn    *websocket.Conn
+	mtx       sync.Mutex
+}
+
+// Listen for events concerning the given uuids. When an event occurs
+// (and after connecting/reconnecting to the event stream), send each
+// uuid to ch. If a {ch, uuid} pair is subscribed twice, the uuid will
+// be sent only once for each update, but two Unsubscribe calls will
+// be needed to stop sending them.
+func (client *arvadosClient) Subscribe(ch chan<- eventMessage, uuid string) {
+	client.mtx.Lock()
+	defer client.mtx.Unlock()
+	if client.notifying == nil {
+		client.notifying = map[string]map[chan<- eventMessage]int{}
+		client.wantClose = make(chan struct{})
+		go client.runNotifier()
+	}
+	chmap := client.notifying[uuid]
+	if chmap == nil {
+		chmap = map[chan<- eventMessage]int{}
+		client.notifying[uuid] = chmap
+	}
+	needSub := true
+	for _, nch := range chmap {
+		if nch > 0 {
+			needSub = false
+			break
+		}
+	}
+	chmap[ch]++
+	if needSub && client.wsconn != nil {
+		go json.NewEncoder(client.wsconn).Encode(map[string]interface{}{
+			"method": "subscribe",
+			"filters": [][]interface{}{
+				{"object_uuid", "=", uuid},
+				{"event_type", "in", []string{"stderr", "crunch-run", "update"}},
+			},
+		})
+	}
+}
+
+func (client *arvadosClient) Unsubscribe(ch chan<- eventMessage, uuid string) {
+	client.mtx.Lock()
+	defer client.mtx.Unlock()
+	chmap := client.notifying[uuid]
+	if n := chmap[ch] - 1; n == 0 {
+		delete(chmap, ch)
+		if len(chmap) == 0 {
+			delete(client.notifying, uuid)
+		}
+		if client.wsconn != nil {
+			go json.NewEncoder(client.wsconn).Encode(map[string]interface{}{
+				"method": "unsubscribe",
+				"filters": [][]interface{}{
+					{"object_uuid", "=", uuid},
+					{"event_type", "in", []string{"stderr", "crunch-run", "update"}},
+				},
+			})
+		}
+	} else if n > 0 {
+		chmap[ch] = n
+	}
+}
+
+func (client *arvadosClient) Close() {
+	client.mtx.Lock()
+	defer client.mtx.Unlock()
+	if client.notifying != nil {
+		client.notifying = nil
+		close(client.wantClose)
+	}
+}
+
+func (client *arvadosClient) runNotifier() {
+reconnect:
+	for {
+		var cluster arvados.Cluster
+		err := client.RequestAndDecode(&cluster, "GET", arvados.EndpointConfigGet.Path, nil, nil)
+		if err != nil {
+			log.Warnf("error getting cluster config: %s", err)
+			time.Sleep(5 * time.Second)
+			continue reconnect
+		}
+		wsURL := cluster.Services.Websocket.ExternalURL
+		wsURL.Scheme = strings.Replace(wsURL.Scheme, "http", "ws", 1)
+		wsURL.Path = "/websocket"
+		wsURLNoToken := wsURL.String()
+		wsURL.RawQuery = url.Values{"api_token": []string{client.AuthToken}}.Encode()
+		conn, err := websocket.Dial(wsURL.String(), "", cluster.Services.Controller.ExternalURL.String())
+		if err != nil {
+			log.Warnf("websocket connection error: %s", err)
+			time.Sleep(5 * time.Second)
+			continue reconnect
+		}
+		log.Printf("connected to websocket at %s", wsURLNoToken)
+
+		client.mtx.Lock()
+		client.wsconn = conn
+		resubscribe := make([]string, 0, len(client.notifying))
+		for uuid := range client.notifying {
+			resubscribe = append(resubscribe, uuid)
+		}
+		client.mtx.Unlock()
+
+		go func() {
+			w := json.NewEncoder(conn)
+			for _, uuid := range resubscribe {
+				w.Encode(map[string]interface{}{
+					"method": "subscribe",
+					"filters": [][]interface{}{
+						{"object_uuid", "=", uuid},
+						{"event_type", "in", []string{"stderr", "crunch-run", "crunchstat", "update"}},
+					},
+				})
+			}
+		}()
+
+		r := json.NewDecoder(conn)
+		for {
+			var msg eventMessage
+			err := r.Decode(&msg)
+			select {
+			case <-client.wantClose:
+				return
+			default:
+				if err != nil {
+					log.Printf("error decoding websocket message: %s", err)
+					client.mtx.Lock()
+					client.wsconn = nil
+					client.mtx.Unlock()
+					go conn.Close()
+					continue reconnect
+				}
+				client.mtx.Lock()
+				for ch := range client.notifying[msg.ObjectUUID] {
+					ch <- msg
+				}
+				client.mtx.Unlock()
+			}
+		}
+	}
+}
+
 type arvadosContainerRunner struct {
 	Client      *arvados.Client
 	Name        string
@@ -22,31 +184,31 @@ type arvadosContainerRunner struct {
 	RAM         int64
 	Prog        string // if empty, run /proc/self/exe
 	Args        []string
-	Mounts      map[string]string
+	Mounts      map[string]map[string]interface{}
+	Priority    int
 }
 
-var (
-	collectionInPathRe = regexp.MustCompile(`^(.*/)?([0-9a-f]{32}\+[0-9]+|[0-9a-z]{5}-[0-9a-z]{5}-[0-9a-z]{15})(/.*)?$`)
-)
-
-func (runner *arvadosContainerRunner) Run() error {
+func (runner *arvadosContainerRunner) Run() (string, error) {
 	if runner.ProjectUUID == "" {
-		return errors.New("cannot run arvados container: ProjectUUID not provided")
+		return "", errors.New("cannot run arvados container: ProjectUUID not provided")
 	}
+
 	mounts := map[string]map[string]interface{}{
 		"/mnt/output": {
-			"kind":     "tmp",
+			"kind":     "collection",
 			"writable": true,
-			"capacity": 100000000000,
 		},
 	}
+	for path, mnt := range runner.Mounts {
+		mounts[path] = mnt
+	}
 
 	prog := runner.Prog
 	if prog == "" {
 		prog = "/mnt/cmd/lightning"
 		cmdUUID, err := runner.makeCommandCollection()
 		if err != nil {
-			return err
+			return "", err
 		}
 		mounts["/mnt/cmd"] = map[string]interface{}{
 			"kind": "collection",
@@ -55,11 +217,9 @@ func (runner *arvadosContainerRunner) Run() error {
 	}
 	command := append([]string{prog}, runner.Args...)
 
-	for uuid, mnt := range runner.Mounts {
-		mounts[mnt] = map[string]interface{}{
-			"kind": "collection",
-			"uuid": uuid,
-		}
+	priority := runner.Priority
+	if priority < 1 {
+		priority = 500
 	}
 	rc := arvados.RuntimeConstraints{
 		VCPUs:        runner.VCPUs,
@@ -77,20 +237,105 @@ func (runner *arvadosContainerRunner) Run() error {
 			"use_existing":        true,
 			"output_path":         "/mnt/output",
 			"runtime_constraints": rc,
-			"priority":            1,
+			"priority":            runner.Priority,
 			"state":               arvados.ContainerRequestStateCommitted,
 		},
 	})
-	log.Print(cr.UUID)
-	return err
+	if err != nil {
+		return "", err
+	}
+	log.Printf("container request UUID: %s", cr.UUID)
+	log.Printf("container UUID: %s", cr.ContainerUUID)
+
+	logch := make(chan eventMessage)
+	client := arvadosClient{Client: runner.Client}
+	defer client.Close()
+	subscribedUUID := ""
+	defer func() {
+		if subscribedUUID != "" {
+			client.Unsubscribe(logch, subscribedUUID)
+		}
+	}()
+
+	ticker := time.NewTicker(5 * time.Second)
+	defer ticker.Stop()
+
+	neednewline := ""
+
+	lastState := cr.State
+	refreshCR := func() {
+		err = runner.Client.RequestAndDecode(&cr, "GET", "arvados/v1/container_requests/"+cr.UUID, nil, nil)
+		if err != nil {
+			fmt.Fprint(os.Stderr, neednewline)
+			log.Printf("error getting container request: %s", err)
+			return
+		}
+		if lastState != cr.State {
+			fmt.Fprint(os.Stderr, neednewline)
+			log.Printf("container request state: %s", cr.State)
+			lastState = cr.State
+		}
+		if subscribedUUID != cr.ContainerUUID {
+			fmt.Fprint(os.Stderr, neednewline)
+			neednewline = ""
+			if subscribedUUID != "" {
+				client.Unsubscribe(logch, subscribedUUID)
+			}
+			client.Subscribe(logch, cr.ContainerUUID)
+			subscribedUUID = cr.ContainerUUID
+		}
+	}
+
+	var reCrunchstat = regexp.MustCompile(`mem .* rss`)
+	for cr.State != arvados.ContainerRequestStateFinal {
+		select {
+		case <-ticker.C:
+			refreshCR()
+		case msg := <-logch:
+			switch msg.EventType {
+			case "update":
+				refreshCR()
+			case "stderr":
+				for _, line := range strings.Split(msg.Properties.Text, "\n") {
+					if line != "" {
+						fmt.Fprint(os.Stderr, neednewline)
+						neednewline = ""
+						log.Print(line)
+					}
+				}
+			case "crunchstat":
+				for _, line := range strings.Split(msg.Properties.Text, "\n") {
+					mem := reCrunchstat.FindString(line)
+					if mem != "" {
+						fmt.Fprintf(os.Stderr, "%s               \r", mem)
+						neednewline = "\n"
+					}
+				}
+			}
+		}
+	}
+	fmt.Fprint(os.Stderr, neednewline)
+
+	var c arvados.Container
+	err = runner.Client.RequestAndDecode(&c, "GET", "arvados/v1/containers/"+cr.ContainerUUID, nil, nil)
+	if err != nil {
+		return "", err
+	} else if c.State != arvados.ContainerStateComplete {
+		return "", fmt.Errorf("container did not complete: %s", c.State)
+	} else if c.ExitCode != 0 {
+		return "", fmt.Errorf("container exited %d", c.ExitCode)
+	}
+	return cr.OutputUUID, err
 }
 
+var collectionInPathRe = regexp.MustCompile(`^(.*/)?([0-9a-f]{32}\+[0-9]+|[0-9a-z]{5}-[0-9a-z]{5}-[0-9a-z]{15})(/.*)?$`)
+
 func (runner *arvadosContainerRunner) TranslatePaths(paths ...*string) error {
 	if runner.Mounts == nil {
-		runner.Mounts = make(map[string]string)
+		runner.Mounts = make(map[string]map[string]interface{})
 	}
 	for _, path := range paths {
-		if *path == "" {
+		if *path == "" || *path == "-" {
 			continue
 		}
 		m := collectionInPathRe.FindStringSubmatch(*path)
@@ -98,12 +343,15 @@ func (runner *arvadosContainerRunner) TranslatePaths(paths ...*string) error {
 			return fmt.Errorf("cannot find uuid in path: %q", *path)
 		}
 		uuid := m[2]
-		mnt, ok := runner.Mounts[uuid]
+		mnt, ok := runner.Mounts["/mnt/"+uuid]
 		if !ok {
-			mnt = "/mnt/" + uuid
-			runner.Mounts[uuid] = mnt
+			mnt = map[string]interface{}{
+				"kind": "collection",
+				"uuid": uuid,
+			}
+			runner.Mounts["/mnt/"+uuid] = mnt
 		}
-		*path = mnt + m[3]
+		*path = "/mnt/" + uuid + m[3]
 	}
 	return nil
 }