Initial distributor implementation

appengine/cmd/dm/distributor/pubsub.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 distributor
+
+import (
+        "encoding/json"
+        "fmt"
+        "net/http"
+        "strconv"
+        "time"
+
+        "github.com/julienschmidt/httprouter"
+        "github.com/luci/luci-go/appengine/tumble"
+        "github.com/luci/luci-go/common/api/dm/service/v1"
+        "github.com/luci/luci-go/common/logging"
+        "github.com/luci/luci-go/server/tokens"
+        "golang.org/x/net/context"
+)
+
+const notifyTopicSuffix = "dm-distributor-notify"
+
+func pubsubTopic(c context.Context) string {

[iannucci, 2016/06/08 02:54:24]

note that all of this is roughly untested. This will be fixed at the beginning
of the actual swarming integration changes.

+        return fmt.Sprintf("projects/%s/topics/%s", getTrimmedAppID(c), notifyTopicSuffix)

[dnj (Google), 2016/06/09 18:00:55]

I would recommend using "common/gcloud/pubsub" here.

You can then call:
pubsub.NewTopic()

It's typed, you can validate, etc.

[iannucci, 2016/06/15 00:46:00]

Done.

+}
+
+// PubsubReciever is the http handler that processes incoming pubsub events

[dnj (Google), 2016/06/09 18:00:55]

nit: HTTP

[iannucci, 2016/06/15 00:46:00]

Done.

+// delivered to topics prepared with TaskDescription.PrepareTopic, and routes
+// them to the appropriate distributor implementation's HandleNotification
+// method.
+//
+// It requires that a Registry be installed in c via WithRegistry.
+func PubsubReciever(c context.Context, rw http.ResponseWriter, r *http.Request, p httprouter.Params) {

[dnj (Google), 2016/06/09 18:00:55]

nit: it might be cleaner to wrap a more specific method:

func handlePubSub(c context.Context, body []byte) (ack bool) {...}

[iannucci, 2016/06/15 00:46:00]

meh.... all this handler stuff is already function soup. I'm going to err on the
'less indirection' side here.

+        defer r.Body.Close()
+
+        type PubsubMessage struct {
+                Attributes map[string]string `json:"attributes"`
+                Data       []byte            `json:"data"`
+                MessageID  string            `json:"message_id"`
+        }
+        type PubsubPushMessage struct {
+                Message      PubsubMessage `json:"message"`
+                Subscription string        `json:"subscription"`
+        }
+        psm := &PubsubPushMessage{}
+
+        err := json.NewDecoder(r.Body).Decode(psm)

[dnj (Google), 2016/06/09 18:00:55]

Why use JSON here and not protobuf?

[iannucci, 2016/06/15 00:46:00]

this is just the pubsub message envelope thingy

+        if err != nil {

[dnj (Google), 2016/06/09 18:00:55]

nit: if err := json.NewDecoder...

[iannucci, 2016/06/15 00:46:00]

Done.

+                logging.Fields{"error": err}.Errorf(c, "Failed to parse pubsub message")
+                http.Error(rw, "Failed to parse pubsub message", http.StatusInternalServerError)
+                return
+        }
+
+        eid, cfgName, err := decodeAuthToken(c, psm.Message.Attributes["auth_token"])
+        if err != nil {
+                logging.WithError(err).Errorf(c, "bad auth_token")
+                // Acknowledge this message, since it'll never be valid.
+                rw.WriteHeader(http.StatusNoContent)
+                return
+        }
+
+        // remove "auth_token" from Attributes to avoid having it pass to the
+        // distributor.
+        delete(psm.Message.Attributes, "auth_token")
+
+        err = tumble.RunMutation(c, &NotifyExecution{
+                cfgName, &Notification{eid, psm.Message.Data, psm.Message.Attributes},
+        })
+        if err != nil {
+                // TODO(riannucci): distinguish between transient/non-transient failures.
+                logging.WithError(err).Errorf(c, "failed to NotifyExecution")
+                rw.WriteHeader(http.StatusInternalServerError)
+                return
+        }
+
+        rw.WriteHeader(http.StatusNoContent)
+}
+
+// pubsubAuthToken describes how to generate HMAC protected tokens used to
+// authenticate PubSub messages.
+var pubsubAuthToken = tokens.TokenKind{
+        Algo:       tokens.TokenAlgoHmacSHA256,
+        Expiration: 48 * time.Hour,
+        SecretKey:  "pubsub_auth_token",
+        Version:    1,
+}
+
+func encodeAuthToken(c context.Context, eid *dm.Execution_ID, cfgName string) (string, error) {
+        return pubsubAuthToken.Generate(c, nil, map[string]string{
+                "quest":     eid.Quest,
+                "attempt":   strconv.FormatUint(uint64(eid.Attempt), 10),
+                "execution": strconv.FormatUint(uint64(eid.Id), 10),
+                "cfgName":   cfgName,
+        }, 0)
+}
+
+func decodeAuthToken(c context.Context, authToken string) (eid *dm.Execution_ID, cfgName string, err error) {
+        items, err := pubsubAuthToken.Validate(c, authToken, nil)
+        if err != nil {
+                return
+        }
+        quest, qok := items["quest"]
+        attempt, aok := items["attempt"]
+        execution, eok := items["execution"]
+        if !qok || !aok || !eok {
+                err = fmt.Errorf("missing keys: %v", items)
+                return
+        }
+        attemptNum, err := strconv.ParseUint(attempt, 10, 32)
+        if err != nil {
+                return
+        }
+        executionNum, err := strconv.ParseUint(execution, 10, 32)
+        if err != nil {
+                return
+        }
+        eid = dm.NewExecutionID(quest, uint32(attemptNum), uint32(executionNum))
+
+        cfgName, ok := items["cfgName"]
+        if !ok {
+                err = fmt.Errorf("missing config name")
+        }
+
+        return
+}