Use swarming in capture_skps and capture_skps_from_pdfs CT tasks

ct/go/worker_scripts/capture_skps_from_pdfs/main.go

 // Application that downloads PDFs and then captures SKPs from them.
 // TODO(rmistry): Capturing and uploading SKPs has been temporarily disabled due
 // to the comment in https://bugs.chromium.org/p/skia/issues/detail?id=5183#c34
 package main
 
 import (
         "flag"
         "fmt"
         "io"
         "io/ioutil"
         "net/url"
         "os"
+        "path"
         "path/filepath"
+        "strconv"
         "strings"
         "sync"
         "time"
 
         "github.com/skia-dev/glog"
 
         "go.skia.org/infra/ct/go/util"
         "go.skia.org/infra/ct/go/worker_scripts/worker_common"
         "go.skia.org/infra/go/common"
         "go.skia.org/infra/go/httputils"
         skutil "go.skia.org/infra/go/util"
 )
 
 const (
         // The number of goroutines that will run in parallel to download PDFs and capture their SKPs.
         WORKER_POOL_SIZE = 10
 )
 
 var (
-        workerNum      = flag.Int("worker_num", 1, "The number of this CT worker. It will be in the {1..100} range.")
+        startRange     = flag.Int("start_range", 1, "The number this worker will capture SKPs from.")
+        num            = flag.Int("num", 100, "The total number of SKPs to capture starting from the start_range.")
         pagesetType    = flag.String("pageset_type", util.PAGESET_TYPE_PDF_1m, "The type of pagesets to use for this run. Eg: PDF1m.")
         chromiumBuild  = flag.String("chromium_build", "", "The specified chromium build. This value is used to find the pdfium_test binary from Google Storage and while uploading the PDFs and SKPs to Google Storage.")
         runID          = flag.String("run_id", "", "The unique run id (typically requester + timestamp).")
         targetPlatform = flag.String("target_platform", util.PLATFORM_LINUX, "The platform the benchmark will run on (Android / Linux).")
 )
 
 func main() {
         defer common.LogPanic()
         worker_common.Init()
-        if !*worker_common.Local {
-                defer util.CleanTmpDir()
-        }
         defer util.TimeTrack(time.Now(), "Capturing SKPs from PDFs")
         defer glog.Flush()
 
         // Validate required arguments.
         if *runID == "" {
                 glog.Error("Must specify --run_id")
                 return
         }
         if *chromiumBuild == "" {
                 glog.Error("Must specify --chromium_build")
                 return
         }
 
         // Instantiate timeout client for downloading PDFs.
         httpTimeoutClient := httputils.NewTimeoutClient()
         // Instantiate GsUtil object.
         gs, err := util.NewGsUtil(nil)
         if err != nil {
                 glog.Error(err)
                 return
         }
 
         // Download PDF pagesets if they do not exist locally.
-        if err := gs.DownloadWorkerArtifacts(util.PAGESETS_DIR_NAME, *pagesetType, *workerNum); err != nil {
+        pathToPagesets := filepath.Join(util.PagesetsDir, *pagesetType)
+        pagesetsToIndex, err := gs.DownloadSwarmingArtifacts(pathToPagesets, util.PAGESETS_DIR_NAME, *pagesetType, *startRange, *num)
+        if err != nil {
                 glog.Error(err)
                 return
         }
-        pathToPagesets := filepath.Join(util.PagesetsDir, *pagesetType)
+        defer skutil.RemoveAll(pathToPagesets)
 
         // Create the dir that PDFs will be stored in.
         pathToPdfs := filepath.Join(util.PdfsDir, *pagesetType, *chromiumBuild)
         // Delete and remake the local PDFs directory.
         skutil.RemoveAll(pathToPdfs)
         skutil.MkdirAll(pathToPdfs, 0700)
         // Cleanup the dir after the task is done.
         defer skutil.RemoveAll(pathToPdfs)
 
         // Create the dir that SKPs will be stored in.
@@ skipping 50 matching lines @@
         for i := 0; i < WORKER_POOL_SIZE; i++ {
                 // Increment the WaitGroup counter.
                 wg.Add(1)
 
                 // Create and run a goroutine closure that captures SKPs.
                 go func() {
                         // Decrement the WaitGroup counter when the goroutine completes.
                         defer wg.Done()
 
                         for pagesetName := range pagesetRequests {
+                                index := strconv.Itoa(pagesetsToIndex[path.Join(pathToPagesets, pagesetName)])
 
                                 // Read the pageset.
                                 pagesetPath := filepath.Join(pathToPagesets, pagesetName)
                                 decodedPageset, err := util.ReadPageset(pagesetPath)
                                 if err != nil {
                                         glog.Errorf("Could not read %s: %s", pagesetPath, err)
                                         continue
                                 }
 
                                 glog.Infof("===== Processing %s =====", pagesetPath)
@@ skipping 10 matching lines @@
                                 // Add protocol if it is missing from the URL.
                                 if !(strings.HasPrefix(pdfURL, "http://") || strings.HasPrefix(pdfURL, "https://")) {
                                         pdfURL = fmt.Sprintf("http://%s", pdfURL)
                                 }
                                 pdfBase, err := getPdfFileName(pdfURL)
                                 if err != nil {
                                         glog.Errorf("Could not parse the URL %s to get a PDF file name: %s", pdfURL, err)
                                         erroredPDFs = append(erroredPDFs, pdfURL)
                                         continue
                                 }
-                                pdfPath := filepath.Join(pathToPdfs, pdfBase)
+                                pdfDirWithIndex := filepath.Join(pathToPdfs, index)
+                                if err := os.MkdirAll(pdfDirWithIndex, 0700); err != nil {
+                                        glog.Errorf("Could not mkdir %s: %s", pdfDirWithIndex, err)
+                                }
+                                pdfPath := filepath.Join(pdfDirWithIndex, pdfBase)
                                 resp, err := httpTimeoutClient.Get(pdfURL)
                                 if err != nil {
                                         glog.Errorf("Could not GET %s: %s", pdfURL, err)
                                         erroredPDFs = append(erroredPDFs, pdfURL)
                                         continue
                                 }
                                 defer skutil.Close(resp.Body)
                                 out, err := os.Create(pdfPath)
                                 if err != nil {
                                         glog.Errorf("Unable to create file %s: %s", pdfPath, err)
@@ skipping 16 matching lines @@
                                 //pdfiumTestArgs := []string{
                                 //      "--skp", pdfPath,
                                 //}
                                 //if err := util.ExecuteCmd(pdfiumLocalPath, pdfiumTestArgs, []string{}, time.Duration(timeoutSecs)*time.Second, nil, nil); err != nil {
                                 //      glog.Errorf("Could not run pdfium_test on %s: %s", pdfPath, err)
                                 //      erroredSKPs = append(erroredSKPs, pdfBase)
                                 //      continue
                                 //}
                                 //
                                 //// Move generated SKPs into the pathToSKPs directory.
-                                //skps, err := filepath.Glob(path.Join(pathToPdfs, fmt.Sprintf("%s.*.skp", pdfBase)))
+                                //skps, err := filepath.Glob(path.Join(pdfDirWithIndex, fmt.Sprintf("%s.*.skp", pdfBase)))
                                 //if err != nil {
                                 //      glog.Errorf("Found no SKPs for %s: %s", pdfBase, err)
                                 //      erroredSKPs = append(erroredSKPs, pdfBase)
                                 //      continue
                                 //}
                                 //for _, skp := range skps {
                                 //      skpBasename := path.Base(skp)
-                                //      dest := path.Join(pathToSkps, skpBasename)
+                                //      destDir := path.Join(pathToSkps, index)
+                                //  if err := os.MkdirAll(destDir, 0700); err != nil {
+                                //              glog.Errorf("Could not mkdir %s: %s", destDir, err)
+                                //      }
+                                //      dest := path.Join(destDir, skpBasename)
                                 //      if err := os.Rename(skp, dest); err != nil {
                                 //              glog.Errorf("Could not move %s to %s: %s", skp, dest, err)
                                 //              continue
                                 //      }
                                 //}
                         }
                 }()
         }
 
         // Wait for all spawned goroutines to complete.
@@ skipping 19 matching lines @@
         //      glog.Errorf("Could not create any SKP in %s", pathToSkps)
         //      return
         //}
         //
         //// Move and validate all SKP files.
         //if err := util.ValidateSKPs(pathToSkps); err != nil {
         //      glog.Error(err)
         //      return
         //}
 
-        // Write timestamp to the PDFs dir.
-        skutil.LogErr(util.CreateTimestampFile(pathToPdfs))
-        // Write timestamp to the SKPs dir.
-        skutil.LogErr(util.CreateTimestampFile(pathToSkps))
-
         // Upload PDFs dir to Google Storage.
-        if err := gs.UploadWorkerArtifacts(util.PDFS_DIR_NAME, filepath.Join(*pagesetType, *chromiumBuild), *workerNum); err != nil {
+        if err := gs.UploadSwarmingArtifacts(util.PDFS_DIR_NAME, filepath.Join(*pagesetType, *chromiumBuild)); err != nil {
                 glog.Error(err)
                 return
         }
         // Upload SKPs dir to Google Storage.
-        if err := gs.UploadWorkerArtifacts(util.SKPS_DIR_NAME, filepath.Join(*pagesetType, *chromiumBuild), *workerNum); err != nil {
+        if err := gs.UploadSwarmingArtifacts(util.SKPS_DIR_NAME, filepath.Join(*pagesetType, *chromiumBuild)); err != nil {
                 glog.Error(err)
                 return
         }
 
         // Summarize errors.
         if len(erroredPDFs) > 0 {
                 glog.Error("The Following URLs could not be downloaded as PDFs:")
                 for _, erroredPDF := range erroredPDFs {
                         glog.Errorf("\t%s", erroredPDF)
                 }
@@ skipping 13 matching lines @@
 //   http://www.ada.gov/emerprepguideprt.pdf will become
 //   www.ada.gov__emerprepguideprt.pdf
 func getPdfFileName(u string) (string, error) {
         p, err := url.Parse(u)
         if err != nil {
                 return "", err
         }
         pdfFileName := fmt.Sprintf("%s%s", p.Host, strings.Replace(p.Path, "/", "__", -1))
         return pdfFileName, nil
 }