[CT] Add ability to download to/upload from swarming GS dir with numerical subdirs

ct/go/util/gs.go

 // Google Storage utility that contains methods for both CT master and worker
 // scripts.
 package util
 
 import (
         "fmt"
         "io"
         "io/ioutil"
         "net/http"
         "os"
+        "path"
         "path/filepath"
+        "strconv"
         "strings"
         "sync"
         "time"
 
         "github.com/skia-dev/glog"
         "go.skia.org/infra/go/auth"
         "go.skia.org/infra/go/gs"
         "go.skia.org/infra/go/util"
         googleapi "google.golang.org/api/googleapi"
         storage "google.golang.org/api/storage/v1"
@@ skipping 199 matching lines @@
 
         if equal, _ := gs.AreTimeStampsEqual(localDir, gsDir); equal {
                 // No need to download artifacts they already exist locally.
                 glog.Infof("Not downloading %s because TIMESTAMPS match", gsDir)
                 return nil
         }
         glog.Infof("Timestamps between %s and %s are different. Downloading from Google Storage", localDir, gsDir)
         return gs.downloadRemoteDir(localDir, gsDir)
 }
 
-func (gs *GsUtil) deleteRemoteDir(gsDir string) error {
+func (gs *GsUtil) DeleteRemoteDir(gsDir string) error {
         // The channel where the GS filepaths to be deleted will be sent to.
         chFilePaths := make(chan string, MAX_CHANNEL_SIZE)
         req := gs.service.Objects.List(GSBucketName).Prefix(gsDir + "/")
         for req != nil {
                 resp, err := req.Do()
                 if err != nil {
                         return fmt.Errorf("Error occured while listing %s: %s", gsDir, err)
                 }
                 for _, result := range resp.Items {
                         chFilePaths <- result.Name
@@ skipping 70 matching lines @@
         gsDir := filepath.Join(dirName, pagesetType, fmt.Sprintf("slave%d", workerNum))
 
         if equal, _ := gs.AreTimeStampsEqual(localDir, gsDir); equal {
                 glog.Infof("Not uploading %s because TIMESTAMPS match", localDir)
                 return nil
         }
         glog.Infof("Timestamps between %s and %s are different. Uploading to Google Storage", localDir, gsDir)
         return gs.UploadDir(localDir, gsDir, true)
 }
 
+// UploadSwarmingArtifact uploads the specified local artifacts to Google Storage.
+func (gs *GsUtil) UploadSwarmingArtifacts(dirName, pagesetType string) error {
+        localDir := path.Join(StorageDir, dirName, pagesetType)
+        gsDir := path.Join(SWARMING_DIR_NAME, dirName, pagesetType)
+
+        return gs.UploadDir(localDir, gsDir, false)
+}
+
+// DownloadSwarmingArtifacts downloads the specified artifacts from Google Storage to a local dir.
+// The Google storage directory is assumed to have numerical subdirs Eg: {1..1000}. This function
+// downloads the contents of those directories into a local directory without the numerical

[dogben, 2016/05/18 15:22:01]

Aren't there files with the same name in multiple subdirs?

[rmistry, 2016/05/19 11:54:10]

Possible but unlikely. If they exist right now they overwrite each other anyway
in the current architecture. They will still overwrite each other in the new
architecture but it will be even less likely due to each bot dealing with fewer
artifacts than before.

+// subdirs.
+// Returns the ranking/index of the downloaded artifact.
+func (gs *GsUtil) DownloadSwarmingArtifacts(localDir, remoteDirName, pagesetType string, startRange, num int) (map[string]int, error) {
+        // Empty the local dir.
+        util.RemoveAll(localDir)
+        // Create the local dir.
+        util.MkdirAll(localDir, 0700)
+
+        gsDir := filepath.Join(SWARMING_DIR_NAME, remoteDirName, pagesetType)
+        endRange := num + startRange - 1
+        // The channel where remote files to be downloaded will be sent to.
+        chRemoteDirs := make(chan string, num)
+        for i := startRange; i <= endRange; i++ {
+                chRemoteDirs <- filepath.Join(gsDir, strconv.Itoa(i))
+        }
+        close(chRemoteDirs)
+
+        // Dictionary of artifacts to its rank/index.
+        artifactToIndex := map[string]int{}
+        // Mutex to control access to the above dictionary.
+        var mtx sync.Mutex
+        // Kick off goroutines to download artifacts and populate the artifactToIndex dictionary.
+        var wg sync.WaitGroup
+        for i := 0; i < GOROUTINE_POOL_SIZE; i++ {
+                wg.Add(1)
+                go func(goroutineNum int) {
+                        defer wg.Done()
+                        for remoteDir := range chRemoteDirs {
+
+                                req := gs.service.Objects.List(GSBucketName).Prefix(remoteDir + "/")
+                                for req != nil {
+                                        resp, err := req.Do()
+                                        if err != nil {
+                                                glog.Errorf("Error occured while listing %s: %s", gsDir, err)
+                                                return
+                                        }
+                                        for _, result := range resp.Items {
+                                                fileName := filepath.Base(result.Name)
+                                                fileGsDir := filepath.Dir(result.Name)
+                                                index, err := strconv.Atoi(path.Base(fileGsDir))
+                                                if err != nil {
+                                                        glog.Errorf("%s was not in expected format: %s", fileGsDir, err)
+                                                }

[dogben, 2016/05/18 15:22:01]

Missing return?

[rmistry, 2016/05/19 11:54:10]

Oops. Done.

+                                                respBody, err := getRespBody(result, gs.client)
+                                                if err != nil {
+                                                        glog.Errorf("Could not fetch %s: %s", result.MediaLink, err)
+                                                        return
+                                                }
+                                                defer util.Close(respBody)

[dogben, 2016/05/18 15:22:01]

If there are many items in chRemoteDirs, deferring the util.Close will prevent
GCing all the respBody's, which could exhaust memory.

I think it would be good to split the goroutine into a different function, since
this function is quite long.

[rmistry, 2016/05/19 11:54:10]

Done.

+                                                outputFile := filepath.Join(localDir, fileName)
+                                                out, err := os.Create(outputFile)
+                                                if err != nil {
+                                                        glog.Errorf("Unable to create file %s: %s", outputFile, err)
+                                                        return
+                                                }
+                                                defer util.Close(out)

[dogben, 2016/05/18 15:22:01]

nit: util.Close will ignore errors on closing the file. Probably doesn't matter,
since it would be errors like the disk is removed before finishing writing.

[rmistry, 2016/05/19 11:54:11]

Yea, I figured just logging an error was enough here.

+                                                if _, err = io.Copy(out, respBody); err != nil {
+                                                        glog.Error(err)
+                                                        return
+                                                }
+                                                glog.Infof("Downloaded gs://%s/%s to %s with goroutine#%d", GSBucketName, result.Name, outputFile, goroutineNum)
+                                                // Sleep for a second after downloading file to avoid bombarding Cloud
+                                                // storage.
+                                                time.Sleep(time.Second)
+                                                mtx.Lock()
+                                                artifactToIndex[path.Join(localDir, fileName)] = index
+                                                mtx.Unlock()
+                                        }
+                                        if len(resp.NextPageToken) > 0 {
+                                                req.PageToken(resp.NextPageToken)
+                                        } else {
+                                                req = nil
+                                        }
+                                }
+                        }
+                }(i + 1)
+        }
+        wg.Wait()
+        return artifactToIndex, nil

[dogben, 2016/05/18 15:22:01]

nit: return error if chRemoteDirs is not empty?

[rmistry, 2016/05/19 11:54:11]

Done.

+}
+
 // UploadDir uploads the specified local dir into the specified Google Storage dir.
 func (gs *GsUtil) UploadDir(localDir, gsDir string, cleanDir bool) error {
         if cleanDir {
                 // Empty the remote dir before uploading to it.
-                util.LogErr(gs.deleteRemoteDir(gsDir))
+                util.LogErr(gs.DeleteRemoteDir(gsDir))
         }
 
         // Construct a dictionary of file paths to their file infos.
         pathsToFileInfos := map[string]os.FileInfo{}
         visit := func(path string, f os.FileInfo, err error) error {
                 if f.IsDir() {
                         return nil
                 }
                 pathsToFileInfos[path] = f
                 return nil
@@ skipping 32 matching lines @@
                                 }
                                 // Sleep for a second after uploading file to avoid bombarding Cloud
                                 // storage.
                                 time.Sleep(time.Second)
                         }
                 }(i + 1)
         }
         wg.Wait()
         return nil
 }