X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/08d080d5a4edccaa579f9a5b5cfa18c9ac471430..c6df16d2af30e989bcfb04f6ef730cde658a9dc9:/services/crunch-run/upload.go

diff --git a/services/crunch-run/upload.go b/services/crunch-run/upload.go
index d50ea581b3..a068a2a77b 100644
--- a/services/crunch-run/upload.go
+++ b/services/crunch-run/upload.go
@@ -17,7 +17,11 @@ import (
 	"git.curoverse.com/arvados.git/sdk/go/keepclient"
 	"git.curoverse.com/arvados.git/sdk/go/manifest"
 	"io"
+	"log"
+	"os"
+	"path/filepath"
 	"strings"
+	"sync"
 )
 
 // Block is a data block in a manifest stream
@@ -77,6 +81,12 @@ func (m *CollectionFileWriter) Close() error {
 	return nil
 }
 
+func (m *CollectionFileWriter) NewFile(fn string) {
+	m.offset += m.length
+	m.length = 0
+	m.fn = fn
+}
+
 func (m *CollectionFileWriter) goUpload() {
 	var errors []error
 	uploader := m.uploader
@@ -93,11 +103,12 @@ func (m *CollectionFileWriter) goUpload() {
 	finish <- errors
 }
 
-// CollectionWriter makes implements creating new Keep collections by opening files
+// CollectionWriter implements creating new Keep collections by opening files
 // and writing to them.
 type CollectionWriter struct {
 	IKeepClient
 	Streams []*CollectionFileWriter
+	mtx     sync.Mutex
 }
 
 // Open a new file for writing in the Keep collection.
@@ -125,6 +136,8 @@ func (m *CollectionWriter) Open(path string) io.WriteCloser {
 		fn}
 	go fw.goUpload()
 
+	m.mtx.Lock()
+	defer m.mtx.Unlock()
 	m.Streams = append(m.Streams, fw)
 
 	return fw
@@ -133,6 +146,9 @@ func (m *CollectionWriter) Open(path string) io.WriteCloser {
 // Finish writing the collection, wait for all blocks to complete uploading.
 func (m *CollectionWriter) Finish() error {
 	var errstring string
+	m.mtx.Lock()
+	defer m.mtx.Unlock()
+
 	for _, stream := range m.Streams {
 		if stream.uploader == nil {
 			continue
@@ -168,7 +184,12 @@ func (m *CollectionWriter) ManifestText() (mt string, err error) {
 
 	var buf bytes.Buffer
 
+	m.mtx.Lock()
+	defer m.mtx.Unlock()
 	for _, v := range m.Streams {
+		if len(v.FileStreamSegments) == 0 {
+			continue
+		}
 		k := v.StreamName
 		if k == "." {
 			buf.WriteString(".")
@@ -177,9 +198,13 @@ func (m *CollectionWriter) ManifestText() (mt string, err error) {
 			k = strings.Replace(k, "\n", "", -1)
 			buf.WriteString("./" + k)
 		}
-		for _, b := range v.Blocks {
-			buf.WriteString(" ")
-			buf.WriteString(b)
+		if len(v.Blocks) > 0 {
+			for _, b := range v.Blocks {
+				buf.WriteString(" ")
+				buf.WriteString(b)
+			}
+		} else {
+			buf.WriteString(" d41d8cd98f00b204e9800998ecf8427e+0")
 		}
 		for _, f := range v.FileStreamSegments {
 			buf.WriteString(" ")
@@ -192,14 +217,24 @@ func (m *CollectionWriter) ManifestText() (mt string, err error) {
 	return buf.String(), nil
 }
 
-func (m *CollectionWriter) WalkFunc(path string, info os.FileInfo, err error, stripPrefix string) error {
+type WalkUpload struct {
+	kc          IKeepClient
+	stripPrefix string
+	streamMap   map[string]*CollectionFileWriter
+	status      *log.Logger
+}
+
+// WalkFunc walks a directory tree, uploads each file found and adds it to the
+// CollectionWriter.
+func (m *WalkUpload) WalkFunc(path string, info os.FileInfo, err error) error {
+
 	if info.IsDir() {
 		return nil
 	}
 
 	var dir string
-	if len(path) > (len(stripPrefix) + len(info.Name()) + 1) {
-		dir = path[len(stripPrefix)+1 : (len(path) - len(info.Name()) - 1)]
+	if len(path) > (len(m.stripPrefix) + len(info.Name()) + 1) {
+		dir = path[len(m.stripPrefix)+1 : (len(path) - len(info.Name()) - 1)]
 	}
 	if dir == "" {
 		dir = "."
@@ -207,9 +242,9 @@ func (m *CollectionWriter) WalkFunc(path string, info os.FileInfo, err error, st
 
 	fn := path[(len(path) - len(info.Name())):]
 
-	if m.Streams[dir] == nil {
-		m.Streams[dir] = &CollectionFileWriter{
-			m.IKeepClient,
+	if m.streamMap[dir] == nil {
+		m.streamMap[dir] = &CollectionFileWriter{
+			m.kc,
 			&manifest.ManifestStream{StreamName: dir},
 			0,
 			0,
@@ -217,48 +252,47 @@ func (m *CollectionWriter) WalkFunc(path string, info os.FileInfo, err error, st
 			make(chan *Block),
 			make(chan []error),
 			""}
-		go m.Streams[dir].goUpload()
+		go m.streamMap[dir].goUpload()
 	}
 
-	stream := m.Streams[dir]
+	fileWriter := m.streamMap[dir]
 
-	fileStart := stream.offset
+	// Reset the CollectionFileWriter for a new file
+	fileWriter.NewFile(fn)
 
 	file, err := os.Open(path)
 	if err != nil {
 		return err
 	}
+	defer file.Close()
 
-	log.Printf("Uploading %v/%v (%v bytes)", dir, fn, info.Size())
+	m.status.Printf("Uploading %v/%v (%v bytes)", dir, fn, info.Size())
 
-	var count int64
-	count, err = io.Copy(stream, file)
+	_, err = io.Copy(fileWriter, file)
 	if err != nil {
 		return err
 	}
 
-	stream.offset += count
-
-	stream.ManifestStream.FileStreamSegments = append(stream.ManifestStream.FileStreamSegments,
-		manifest.FileStreamSegment{uint64(fileStart), uint64(count), fn})
+	// Commits the current file.  Legal to call this repeatedly.
+	fileWriter.Close()
 
 	return nil
 }
 
-func WriteTree(kc IKeepClient, root string) (manifest string, err error) {
-	mw := CollectionWriter{kc, root, map[string]*ManifestStreamWriter{}}
-	err = filepath.Walk(root, func(path string, info os.FileInfo, err error) {
-		return mw.WalkFunc(path, info, err, root)
-	})
+func (cw *CollectionWriter) WriteTree(root string, status *log.Logger) (manifest string, err error) {
+	streamMap := make(map[string]*CollectionFileWriter)
+	wu := &WalkUpload{cw.IKeepClient, root, streamMap, status}
+	err = filepath.Walk(root, wu.WalkFunc)
 
 	if err != nil {
 		return "", err
 	}
 
-	err = mw.Finish()
-	if err != nil {
-		return "", err
+	cw.mtx.Lock()
+	for _, st := range streamMap {
+		cw.Streams = append(cw.Streams, st)
 	}
+	cw.mtx.Unlock()
 
-	return mw.ManifestText(), nil
+	return cw.ManifestText()
 }