isolate: pull exparchive filewalk code into its own func

client/cmd/isolate/exp_archive.go

 // Copyright 2015 The LUCI Authors. All rights reserved.
 // Use of this source code is governed under the Apache License, Version 2.0
 // that can be found in the LICENSE file.
 
 package main
 
 import (
         "encoding/json"
         "errors"
         "fmt"
@@ skipping 59 matching lines @@
 // Item represents a file or symlink referenced by an isolate file.
 type Item struct {
         Path    string
         RelPath string
         Size    int64
         Mode    os.FileMode
 
         Digest isolated.HexDigest
 }
 
+// itemGroup is a list of Items, plus a count of the aggregate size.
+type itemGroup struct {
+        items     []*Item
+        totalSize int64
+}
+
+func (ig *itemGroup) AddItem(item *Item) {
+        ig.items = append(ig.items, item)
+        ig.totalSize += item.Size
+}
+
+// partitioningWalker contains the state necessary to partition isolate deps by handling multiple os.WalkFunc invocations.
+type partitioningWalker struct {
+        // rootDir must be initialized before walkFn is called.
+        rootDir string
+
+        parts partitionedDeps
+}
+
+// partitionedDeps contains a list of items to be archived, partitioned into symlinks and files categorized by size.
+type partitionedDeps struct {
+        links          itemGroup
+        filesToArchive itemGroup
+        indivFiles     itemGroup
+}
+
+// walkFn implements filepath.WalkFunc, for use traversing a directory hierarchy to be isolated.
+// It accumulates files in pw.parts, partitioned into symlinks and files categorized by size.
+func (pw *partitioningWalker) walkFn(path string, info os.FileInfo, err error) error {
+        if err != nil {
+                return err
+        }
+        if info.IsDir() {
+                return nil
+        }
+
+        relPath, err := filepath.Rel(pw.rootDir, path)
+        if err != nil {
+                return err
+        }
+
+        item := &Item{
+                Path:    path,
+                RelPath: relPath,
+                Mode:    info.Mode(),
+                Size:    info.Size(),
+        }
+
+        switch {
+        case item.Mode&os.ModeSymlink == os.ModeSymlink:
+                pw.parts.links.AddItem(item)
+        case item.Size < archiveThreshold:
+                pw.parts.filesToArchive.AddItem(item)
+        default:
+                pw.parts.indivFiles.AddItem(item)
+        }
+        return nil
+}
+
+// partitionDeps walks each of the deps, partioning the results into symlinks and files categorized by size.
+func partitionDeps(deps []string, rootDir string) (partitionedDeps, error) {
+        walker := partitioningWalker{rootDir: rootDir}
+        for _, dep := range deps {
+                // Try to walk dep. If dep is a file (or symlink), the inner function is called exactly once.
+                if err := filepath.Walk(filepath.Clean(dep), walker.walkFn); err != nil {
+                        return partitionedDeps{}, err
+                }
+        }
+        return walker.parts, nil
+}
+
 // main contains the core logic for experimental archive.
 func (c *expArchiveRun) main() error {
         // TODO(djd): This func is long and has a lot of internal complexity (like,
         // such as, archiveCallback). Refactor.
 
         start := time.Now()
         archiveOpts := &c.isolateFlags.ArchiveOptions
         // Parse the incoming isolate file.
         deps, rootDir, isol, err := isolate.ProcessIsolate(archiveOpts)
         if err != nil {
@@ skipping 12 matching lines @@
         }
         client := isolatedclient.New(nil, authCl, c.isolatedFlags.ServerURL, c.isolatedFlags.Namespace, nil, nil)
 
         // Set up a checker and uploader. We limit the uploader to one concurrent
         // upload, since the uploads are all coming from disk (with the exception of
         // the isolated JSON itself) and we only want a single goroutine reading from
         // disk at once.
         checker := NewChecker(ctx, client)
         uploader := NewUploader(ctx, client, 1)
 
-        // Walk each of the deps, partioning the results into symlinks and files categorised by size.
-        var links, archiveFiles, indivFiles []*Item
-        var archiveSize, indivSize int64 // Cumulative size of archived/individual files.
-        for _, dep := range deps {
-                // Try to walk dep. If dep is a file (or symlink), the inner function is called exactly once.
-                err := filepath.Walk(filepath.Clean(dep), func(path string, info os.FileInfo, err error) error {
-                        if err != nil {
-                                return err
-                        }
-                        if info.IsDir() {
-                                return nil
-                        }
-
-                        relPath, err := filepath.Rel(rootDir, path)
-                        if err != nil {
-                                return err
-                        }
-
-                        item := &Item{
-                                Path:    path,
-                                RelPath: relPath,
-                                Mode:    info.Mode(),
-                                Size:    info.Size(),
-                        }
-
-                        switch {
-                        case item.Mode&os.ModeSymlink == os.ModeSymlink:
-                                links = append(links, item)
-                        case item.Size < archiveThreshold:
-                                archiveFiles = append(archiveFiles, item)
-                                archiveSize += item.Size
-                        default:
-                                indivFiles = append(indivFiles, item)
-                                indivSize += item.Size
-                        }
-                        return nil
-                })
-                if err != nil {
-                        return err
-                }
+        parts, err := partitionDeps(deps, rootDir)
+        if err != nil {
+                return fmt.Errorf("partitioning deps: %v", err)
         }
 
         // Construct a map of the files that constitute the isolate.
         files := make(map[string]isolated.File)
 
-        log.Printf("Isolate expanded to %d files (total size %s) and %d symlinks", len(archiveFiles)+len(indivFiles), humanize.Bytes(uint64(archiveSize+indivSize)), len(links))
-        log.Printf("\t%d files (%s) to be isolated individually", len(indivFiles), humanize.Bytes(uint64(indivSize)))
-        log.Printf("\t%d files (%s) to be isolated in archives", len(archiveFiles), humanize.Bytes(uint64(archiveSize)))
+        numFiles := len(parts.filesToArchive.items) + len(parts.indivFiles.items)
+        filesSize := uint64(parts.filesToArchive.totalSize + parts.indivFiles.totalSize)
+        log.Printf("Isolate expanded to %d files (total size %s) and %d symlinks", numFiles, humanize.Bytes(filesSize), len(parts.links.items))
+        log.Printf("\t%d files (%s) to be isolated individually", len(parts.indivFiles.items), humanize.Bytes(uint64(parts.indivFiles.totalSize)))
+        log.Printf("\t%d files (%s) to be isolated in archives", len(parts.filesToArchive.items), humanize.Bytes(uint64(parts.filesToArchive.totalSize)))
 
         // Handle the symlinks.
-        for _, item := range links {
+        for _, item := range parts.links.items {
                 l, err := os.Readlink(item.Path)
                 if err != nil {
                         return fmt.Errorf("unable to resolve symlink for %q: %v", item.Path, err)
                 }
                 files[item.RelPath] = isolated.SymLink(l)
         }
 
         // Handle the small to-be-archived files.
-        bundles := ShardItems(archiveFiles, archiveMaxSize)
+        bundles := ShardItems(parts.filesToArchive.items, archiveMaxSize)
         log.Printf("\t%d TAR archives to be isolated", len(bundles))
 
         for _, bundle := range bundles {
                 bundle := bundle
                 digest, tarSize, err := bundle.Digest()
                 if err != nil {
                         return err
                 }
 
                 log.Printf("Created tar archive %q (%s)", digest, humanize.Bytes(uint64(tarSize)))
@@ skipping 13 matching lines @@
                                 return
                         }
                         log.Printf("QUEUED %q for upload", item.RelPath)
                         uploader.Upload(item.RelPath, bundle.Contents, ps, func() {
                                 log.Printf("UPLOADED %q", item.RelPath)
                         })
                 })
         }
 
         // Handle the large individually-uploaded files.
-        for _, item := range indivFiles {
+        for _, item := range parts.indivFiles.items {
                 d, err := hashFile(item.Path)
                 if err != nil {
                         return err
                 }
                 item.Digest = d
                 files[item.RelPath] = isolated.BasicFile(item.Digest, int(item.Mode), item.Size)
                 checker.AddItem(item, false, func(item *Item, ps *isolatedclient.PushState) {
                         if ps == nil {
                                 return
                         }
                         log.Printf("QUEUED %q for upload", item.RelPath)
                         uploader.UploadFile(item, ps, func() {
                                 log.Printf("UPLOADED %q", item.RelPath)
                         })
                 })
         }
 
         // Marshal the isolated file into JSON, and create an Item to describe it.
         isol.Files = files
-        isolJSON, err := json.Marshal(isol)
+        var isolJSON []byte
+        isolJSON, err = json.Marshal(isol)
         if err != nil {
                 return err
         }
         isolItem := &Item{
                 Path:    archiveOpts.Isolated,
                 RelPath: filepath.Base(archiveOpts.Isolated),
                 Digest:  isolated.HashBytes(isolJSON),
                 Size:    int64(len(isolJSON)),
         }
 
@@ skipping 96 matching lines @@
 }
 
 func hashFile(path string) (isolated.HexDigest, error) {
         f, err := os.Open(path)
         if err != nil {
                 return "", err
         }
         defer f.Close()
         return isolated.Hash(f)
 }