Add snapshotting for CIPD packages and dimensions to DM.

dm/appengine/distributor/swarming/v1/distributor.go

 // Copyright 2016 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 swarming
 
 import (
         "encoding/json"
         "fmt"
         "net/http"
@@ skipping 27 matching lines @@
 
 var _ distributor.D = (*swarmingDist)(nil)
 
 var swarmBotLookup = map[string]dm.AbnormalFinish_Status{
         "BOT_DIED":  dm.AbnormalFinish_CRASHED,
         "CANCELED":  dm.AbnormalFinish_CANCELLED,
         "EXPIRED":   dm.AbnormalFinish_EXPIRED,
         "TIMED_OUT": dm.AbnormalFinish_TIMED_OUT,
 }
 
+func cipdPackageFromSwarm(pkg *swarm.SwarmingRpcsCipdPackage) *sv1.CipdPackage {
+        return &sv1.CipdPackage{Name: pkg.PackageName, Version: pkg.Version}
+}
+
+func cipdSpecFromSwarm(pkgs *swarm.SwarmingRpcsCipdPins) *sv1.CipdSpec {
+        ret := &sv1.CipdSpec{
+                Client: cipdPackageFromSwarm(pkgs.ClientPackage),
+                ByPath: map[string]*sv1.CipdSpec_CipdPackages{},
+        }
+        for _, pkg := range pkgs.Packages {
+                pkgs, ok := ret.ByPath[pkg.Path]
+                if !ok {
+                        pkgs = &sv1.CipdSpec_CipdPackages{}
+                        ret.ByPath[pkg.Path] = pkgs
+                }
+                pkgs.Pkg = append(pkgs.Pkg, cipdPackageFromSwarm(pkg))
+        }
+        return ret
+}
+
 func toSwarmMap(m map[string]string) []*swarm.SwarmingRpcsStringPair {
         ret := make([]*swarm.SwarmingRpcsStringPair, 0, len(m))
         for key, value := range m {
                 ret = append(ret, &swarm.SwarmingRpcsStringPair{
                         Key: key, Value: value})
         }
         return ret
 }
 
 func httpClients(c context.Context) (anonC, authC *http.Client) {
@@ skipping 42 matching lines @@
 }
 
 func (d *swarmingDist) Run(desc *dm.Quest_Desc, auth *dm.Execution_Auth, prev *dm.JsonResult) (tok distributor.Token, _ time.Duration, err error) {
         id := auth.Id
 
         params, err := parseParams(desc)
         if err != nil {
                 return
         }
 
+        prevParsed := (*sv1.Result)(nil)
+        if prev != nil {
+                prevParsed = &sv1.Result{}
+                if err = jsonpb.UnmarshalString(prev.Object, prevParsed); err != nil {
+                        err = errors.Annotate(err).Reason("parsing previous result").Err()
+                        return
+                }
+        }
+
         isoCtx, _ := context.WithTimeout(d, 30*time.Second)
         iso, err := prepIsolate(isoCtx, d.sCfg.Isolate.Url, desc, auth, prev, params)
         if err != nil {
                 err = errors.Annotate(err).Reason("prepping Isolated").Err()
                 return
         }
 
+        cipdInput := (*swarm.SwarmingRpcsCipdInput)(nil)
+        if prevParsed != nil {
+                cipdInput = prevParsed.CipdPins.ToCipdInput()
+        } else {
+                cipdInput = params.Job.Inputs.Cipd.ToCipdInput()
+        }
+
         topic, token, err := d.cfg.PrepareTopic(d, auth.Id)
         if err != nil {
                 err = errors.Annotate(err).Reason("preparing topic").Err()
                 return
         }
 
-        cipdInput := (*swarm.SwarmingRpcsCipdInput)(nil)
-        if len(params.Job.Inputs.Packages) > 0 {
-                cipdInput := &swarm.SwarmingRpcsCipdInput{
-                        Server: params.Job.Inputs.CipdServer,
-                }
-                for _, pkg := range params.Job.Inputs.Packages {
-                        cipdInput.Packages = append(cipdInput.Packages, &swarm.SwarmingRpcsCipdPackage{
-                                PackageName: pkg.Name,
-                                Path:        pkg.Path,
-                                Version:     pkg.Version,
-                        })
-                }
-        }
-
-        dims := []*swarm.SwarmingRpcsStringPair(nil)
-        for key, value := range params.Scheduling.Dimensions {
-                dims = append(dims, &swarm.SwarmingRpcsStringPair{Key: key, Value: value})
-        }
-
         prefix := params.Meta.NamePrefix
         if len(prefix) > 0 {
                 prefix += " / "
         }
 
+        dims := make(map[string]string, len(params.Scheduling.Dimensions))
+        for k, v := range params.Scheduling.Dimensions {
+                dims[k] = v
+        }
+        if prevParsed != nil {
+                for k, v := range prevParsed.SnapshotDimensions {
+                        dims[k] = v
+                }
+        }
+
         tags := []string{
                 "requestor:DM",
                 "requestor:" + info.TrimmedAppID(d),
                 "requestor:swarming_v1",
                 fmt.Sprintf("quest:%s", id.Quest),
                 fmt.Sprintf("attempt:%s|%d", id.Quest, id.Attempt),
                 fmt.Sprintf("execution:%s|%d|%d", id.Quest, id.Attempt, id.Id),
         }
 
         rslt := (*swarm.SwarmingRpcsTaskRequestMetadata)(nil)
         err = retry.Retry(d, retry.Default, func() (err error) {
                 rpcCtx, _ := context.WithTimeout(d, 10*time.Second)
                 rslt, err = newSwarmClient(rpcCtx, d.sCfg).Tasks.New(&swarm.SwarmingRpcsNewTaskRequest{
                         ExpirationSecs: int64(desc.Meta.Timeouts.Start.Duration().Seconds()),
                         Name:           fmt.Sprintf("%s%s|%d|%d", prefix, id.Quest, id.Attempt, id.Id),
 
-                        // Priority is already pre-Normalize()'d
                         Priority: int64(params.Scheduling.Priority),
 
                         Properties: &swarm.SwarmingRpcsTaskProperties{
                                 CipdInput:            cipdInput,
-                                Dimensions:           toSwarmMap(params.Scheduling.Dimensions),
+                                Dimensions:           toSwarmMap(dims),
                                 Env:                  toSwarmMap(params.Job.Env),
                                 ExecutionTimeoutSecs: int64(desc.Meta.Timeouts.Run.Duration().Seconds()),
                                 GracePeriodSecs:      int64(desc.Meta.Timeouts.Stop.Duration().Seconds()),
                                 IoTimeoutSecs:        int64(params.Scheduling.IoTimeout.Duration().Seconds()),
                                 InputsRef:            iso,
                         },
 
                         PubsubTopic:     topic.String(),
                         PubsubAuthToken: token,
 
@@ skipping 11 matching lines @@
 }
 
 func (d *swarmingDist) Cancel(q *dm.Quest_Desc, tok distributor.Token) error {
         return retry.Retry(d, retry.Default, func() (err error) {
                 ctx, _ := context.WithTimeout(d, 10*time.Second)
                 _, err = newSwarmClient(ctx, d.sCfg).Task.Cancel(string(tok)).Context(ctx).Do()
                 return
         }, retry.LogCallback(d, "swarm.Task.Cancel"))
 }
 
-func (d *swarmingDist) GetStatus(_ *dm.Quest_Desc, tok distributor.Token) (*dm.Result, error) {
+func snapshotDimensions(p *sv1.Parameters, dims []*swarm.SwarmingRpcsStringListPair) map[string]string {
+        if len(p.Scheduling.SnapshotDimensions) == 0 {
+                return nil
+        }
+        allDims := map[string]string{}
+        for _, dim := range dims {
+                allDims[dim.Key] = dim.Value[len(dim.Value)-1]
+        }
+
+        ret := make(map[string]string, len(p.Scheduling.SnapshotDimensions))
+        for _, k := range p.Scheduling.SnapshotDimensions {
+                if v, ok := allDims[k]; ok {
+                        ret[k] = v
+                }
+        }
+        return ret
+}
+
+func (d *swarmingDist) GetStatus(q *dm.Quest_Desc, tok distributor.Token) (*dm.Result, error) {
         rslt := (*swarm.SwarmingRpcsTaskResult)(nil)
 
         err := retry.Retry(d, retry.Default, func() (err error) {
                 ctx, _ := context.WithTimeout(d, 10*time.Second)
                 rslt, err = newSwarmClient(ctx, d.sCfg).Task.Result(string(tok)).Context(ctx).Do()
                 return
         }, retry.LogCallback(d, fmt.Sprintf("swarm.Task.Result(%s)", tok)))
         if err != nil {
                 if gerr := err.(*googleapi.Error); gerr != nil {
                         if gerr.Code == http.StatusNotFound {
@@ skipping 19 matching lines @@
                         ret.AbnormalFinish = &dm.AbnormalFinish{
                                 Status: dm.AbnormalFinish_CRASHED,
                                 Reason: fmt.Sprintf("swarming: COMPLETED/InternalFailure(%d)", rslt.ExitCode),
                         }
                         break
                 }
 
                 retData := &sv1.Result{
                         ExitCode: rslt.ExitCode,
                 }
+                if rslt.CipdPins != nil {
+                        retData.CipdPins = cipdSpecFromSwarm(rslt.CipdPins)
+                }
+                params, err := parseParams(q)
+                if err != nil {
+                        return nil, err
+                }
+                retData.SnapshotDimensions = snapshotDimensions(params, rslt.BotDimensions)
+
                 if ref := rslt.OutputsRef; ref != nil {
                         retData.IsolatedOutdir = &sv1.IsolatedRef{
                                 Id: ref.Isolated, Server: ref.Isolatedserver}
                 }
                 data, err := (&jsonpb.Marshaler{OrigName: true}).MarshalToString(retData)
                 if err != nil {
                         panic(err)
                 }
 
                 ret.Data = &dm.JsonResult{
@@ skipping 51 matching lines @@
 
 func factory(c context.Context, cfg *distributor.Config) (distributor.D, error) {
         return &swarmingDist{c, cfg, cfg.Content.(*sv1.Config)}, nil
 }
 
 // AddFactory adds this distributor implementation into the distributor
 // Registry.
 func AddFactory(m distributor.FactoryMap) {
         m[(*sv1.Config)(nil)] = factory
 }