Isolate: Use generators instead of seekers

client/archiver/archiver.go

 // Copyright 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 package archiver
 
 import (
         "errors"
         "fmt"
         "io"
@@ skipping 57 matching lines @@
         Error() error
         // Digest returns the calculated digest once calculated, empty otherwise.
         Digest() isolated.HexDigest
 }
 
 // Archiver is an high level interface to an isolatedclient.IsolateServer.
 type Archiver interface {
         common.Canceler
         // Push schedules item upload to the isolate server.
         // Smaller priority value means earlier processing.
-        Push(displayName string, src io.ReadSeeker, priority int64) Future
+        Push(displayName string, src isolatedclient.Source, priority int64) Future
         // PushFile schedules file upload to the isolate server.
         // Smaller priority value means earlier processing.
         PushFile(displayName, path string, priority int64) Future
         Stats() *Stats
 }
 
 // UploadStat is the statistic for a single upload.
 type UploadStat struct {
         Duration time.Duration
         Size     units.Size
@@ skipping 94 matching lines @@
 }
 
 // Private details.
 
 // archiverItem is an item to process. Implements Future.
 //
 // It is caried over from pipeline stage to stage to do processing on it.
 type archiverItem struct {
         // Immutable.
         displayName string         // Name to use to qualify this item
-        path        string         // Set when source is a file on disk
         wgHashed    sync.WaitGroup // Released once .digestItem.Digest is set
         priority    int64          // Lower values - earlier hashing and uploading.
+        path        string         // Set when the source is a file on disk.
         a           *archiver
 
         // Mutable.
         lock       sync.Mutex
         err        error                                    // Item specific error
         digestItem isolateservice.HandlersEndpointsV1Digest // Mutated by hashLoop(), used by doContains()
         linked     []*archiverItem                          // Deduplicated item.
 
         // Mutable but not accessible externally.
-        src   io.ReadSeeker             // Source of data
-        state *isolatedclient.PushState // Server-side push state for cache miss
+        source isolatedclient.Source     // Source of data
+        state  *isolatedclient.PushState // Server-side push state for cache miss
 }
 
-func newArchiverItem(a *archiver, displayName, path string, src io.ReadSeeker, priority int64) *archiverItem {
+func newArchiverItem(a *archiver, displayName, path string, source isolatedclient.Source, priority int64) *archiverItem {
         tracer.CounterAdd(a, "itemsProcessing", 1)
-        i := &archiverItem{a: a, displayName: displayName, path: path, src: src, priority: priority}
+        i := &archiverItem{a: a, displayName: displayName, path: path, source: source, priority: priority}
         i.wgHashed.Add(1)
         return i
 }
 
 func (i *archiverItem) Close() error {
         tracer.CounterAdd(i.a, "itemsProcessing", -1)
         i.a = nil
         return nil
 }
 
@@ skipping 10 matching lines @@
         defer i.lock.Unlock()
         return i.err
 }
 
 func (i *archiverItem) Digest() isolated.HexDigest {
         i.lock.Lock()
         defer i.lock.Unlock()
         return isolated.HexDigest(i.digestItem.Digest)
 }
 
+func (i *archiverItem) isFile() bool {
+        return len(i.path) != 0
+}
+
 func (i *archiverItem) setErr(err error) {
         if err == nil {
                 panic("internal error")
         }
         i.lock.Lock()
         defer i.lock.Unlock()
         if i.err == nil {
                 i.err = err
                 for _, child := range i.linked {
                         child.lock.Lock()
                         child.err = err
                         child.lock.Unlock()
                 }
         }
-        // TODO(maruel): Support Close().
-        i.src = nil
 }
 
 func (i *archiverItem) calcDigest() error {
         defer i.wgHashed.Done()
         var d isolateservice.HandlersEndpointsV1Digest
-        if i.path != "" {
-                // Open and hash the file.
-                var err error
-                if d, err = isolated.HashFile(i.path); err != nil {
-                        i.setErr(err)
-                        return fmt.Errorf("hash(%s) failed: %s\n", i.DisplayName(), err)
-                }
-        } else {
-                // Use src instead.
-                h := isolated.GetHash()
-                size, err := io.Copy(h, i.src)
-                if err != nil {
-                        i.setErr(err)
-                        return fmt.Errorf("read(%s) failed: %s\n", i.DisplayName(), err)
-                }
-                if pos, err := i.src.Seek(0, os.SEEK_SET); err != nil || pos != 0 {
-                        err = fmt.Errorf("seek(%s) failed: %s\n", i.DisplayName(), err)
-                        i.setErr(err)
-                        return err
-                }
-                d = isolateservice.HandlersEndpointsV1Digest{Digest: string(isolated.Sum(h)), IsIsolated: true, Size: size}
+
+        src, err := i.source()
+        if err != nil {
+                return fmt.Errorf("source(%s) failed: %s\n", i.DisplayName(), err)
         }
+        defer src.Close()
+
+        h := isolated.GetHash()
+        size, err := io.Copy(h, src)
+        if err != nil {
+                i.setErr(err)
+                return fmt.Errorf("read(%s) failed: %s\n", i.DisplayName(), err)
+        }
+        d = isolateservice.HandlersEndpointsV1Digest{Digest: string(isolated.Sum(h)), IsIsolated: true, Size: size}
+
         i.lock.Lock()
         defer i.lock.Unlock()
         i.digestItem = d
 
         for _, child := range i.linked {
                 child.lock.Lock()
                 child.digestItem = d
                 child.lock.Unlock()
                 child.wgHashed.Done()
         }
         return nil
 }
 
 func (i *archiverItem) link(child *archiverItem) {
         child.lock.Lock()
         defer child.lock.Unlock()
-        if child.src != nil || child.state != nil || child.linked != nil || child.err != nil {
+        if !child.isFile() || child.state != nil || child.linked != nil || child.err != nil {
                 panic("internal error")
         }
         i.lock.Lock()
         defer i.lock.Unlock()
         i.linked = append(i.linked, child)
         if i.digestItem.Digest != "" {
                 child.digestItem = i.digestItem
                 child.err = i.err
                 child.wgHashed.Done()
         } else if i.err != nil {
@@ skipping 61 matching lines @@
 }
 
 func (a *archiver) CancelationReason() error {
         return a.canceler.CancelationReason()
 }
 
 func (a *archiver) Channel() <-chan error {
         return a.canceler.Channel()
 }
 
-func (a *archiver) Push(displayName string, src io.ReadSeeker, priority int64) Future {
-        i := newArchiverItem(a, displayName, "", src, priority)
-        if pos, err := i.src.Seek(0, os.SEEK_SET); pos != 0 || err != nil {
-                i.setErr(fmt.Errorf("seek(%s) failed: %s\n", i.DisplayName(), err))
-                i.wgHashed.Done()
-                return i
-        }
-        return a.push(i)
+func (a *archiver) Push(displayName string, source isolatedclient.Source, priority int64) Future {
+        return a.push(newArchiverItem(a, displayName, "", source, priority))
 }
 
 func (a *archiver) PushFile(displayName, path string, priority int64) Future {
-        return a.push(newArchiverItem(a, displayName, path, nil, priority))
+        source := func() (io.ReadCloser, error) {
+                return os.Open(path)
+        }
+        return a.push(newArchiverItem(a, displayName, path, source, priority))
 }
 
 func (a *archiver) Stats() *Stats {
         a.statsLock.Lock()
         defer a.statsLock.Unlock()
         return a.stats.deepCopy()
 }
 
 func (a *archiver) push(item *archiverItem) Future {
         if a.pushLocked(item) {
@@ skipping 52 matching lines @@
                 // This loop must never block and must be as fast as it can as it is
                 // functionally equivalent to running with a.closeLock held.
                 if err := a.CancelationReason(); err != nil {
                         item.setErr(err)
                         item.Close()
                         item.wgHashed.Done()
                         continue
                 }
                 // TODO(maruel): Resolve symlinks for further deduplication? Depends on the
                 // use case, not sure the trade off is worth.
-                if item.path != "" {
+                if item.isFile() {
                         if previous, ok := seen[item.path]; ok {
                                 previous.link(item)
                                 // TODO(maruel): Semantically weird.
                                 item.Close()
                                 continue
                         }
                 }
 
                 buildUp = append(buildUp, item)
                 seen[item.path] = item
@@ skipping 139 matching lines @@
                         a.progress.Update(groupUpload, groupUploadTodo, 1)
                         a.progress.Update(groupUpload, groupUploadTodoSize, items[index].digestItem.Size)
                         a.stage4UploadChan <- items[index]
                 }
         }
         log.Printf("Looked up %d items\n", len(items))
 }
 
 // doUpload is called by stage 4.
 func (a *archiver) doUpload(item *archiverItem) {
-        var src io.ReadSeeker
-        if item.src == nil {
-                f, err := os.Open(item.path)
-                if err != nil {
-                        a.Cancel(err)
-                        item.setErr(err)
-                        item.Close()
-                        return
-                }
-                defer f.Close()
-                src = f
-        } else {
-                src = item.src
-                item.src = nil
-        }
         start := time.Now()
-        if err := a.is.Push(item.state, src); err != nil {
+        if err := a.is.Push(item.state, item.source); err != nil {
                 err = fmt.Errorf("push(%s) failed: %s\n", item.path, err)
                 a.Cancel(err)
                 item.setErr(err)
         } else {
                 a.progress.Update(groupUpload, groupUploadDone, 1)
                 a.progress.Update(groupUpload, groupUploadDoneSize, item.digestItem.Size)
         }
         item.Close()
         size := units.Size(item.digestItem.Size)
         u := &UploadStat{time.Since(start), size, item.DisplayName()}
         a.statsLock.Lock()
         a.stats.Pushed = append(a.stats.Pushed, u)
         a.statsLock.Unlock()
         log.Printf("Uploaded %7s: %s\n", size, item.DisplayName())
 }