Milo: Move buildbucket pubsub sub from buildbucket project to milo project

milo/common/pubsub.go

 package common
 
 import (
         "encoding/base64"
         "errors"
         "fmt"
         "net/url"
         "strings"
+        "sync"
         "time"
 
         "cloud.google.com/go/pubsub"
         "golang.org/x/net/context"
 
         "github.com/luci/gae/service/info"
+        luciErrors "github.com/luci/luci-go/common/errors"

[iannucci, 2017/07/13 01:54:02]

just use this instead of errors. there's no reason to import both.

[Ryan Tseng, 2017/07/14 00:51:06]

Done.

         "github.com/luci/luci-go/common/logging"
         "github.com/luci/luci-go/milo/api/config"
 )
 
 var pubSubClientKey = "stores a pubsubClient"

[iannucci, 2017/07/13 01:54:02]

*pubsubClients ?

[Ryan Tseng, 2017/07/14 00:51:06]

Done.

 
 type PubSubMessage struct {
         Attributes map[string]string `json:"attributes"`
         Data       string            `json:"data"`
         MessageID  string            `json:"message_id"`
 }
 
 type PubSubSubscription struct {
         Message      PubSubMessage `json:"message"`
         Subscription string        `json:"subscription"`
 }
 
 var errNotExist = errors.New("does not exist")
 
 // GetData returns the expanded form of Data (decoded from base64).
 func (m *PubSubSubscription) GetData() ([]byte, error) {
         return base64.StdEncoding.DecodeString(m.Message.Data)
 }
 
 // pubsubClient is an interface representing a pubsub.Client containing only
 // the functions that Milo calls.  Internal use only, can be swapped
 // out for testing.
 type pubsubClient interface {
         // getTopic returns the pubsub topic if it exists, a notExist error if
         // it does not exist, or an error if there was an error.
-        getTopic(string) (*pubsub.Topic, error)
+        getTopic(context.Context, string) (*pubsub.Topic, error)
 
         // getSubscription returns the pubsub subscription if it exists,
         // a notExist error if it does not exist, or an error if there was an error.
-        getSubscription(string) (*pubsub.Subscription, error)
+        getSubscription(context.Context, string) (*pubsub.Subscription, error)
 
-        createSubscription(string, pubsub.SubscriptionConfig) (
+        createSubscription(context.Context, string, pubsub.SubscriptionConfig) (
                 *pubsub.Subscription, error)
 }
 
+type pubsubClients struct {

[iannucci, 2017/07/13 02:16:27]

now that I'm looking at this again... this is a bit weird. Can we just make the
clients on demand?

[Ryan Tseng, 2017/07/14 00:51:06]

Made this back into a client factory

+        clients map[string]pubsubClient
+        lock    sync.RWMutex
+}
+
+func newClientsBundle() *pubsubClients {
+        return &pubsubClients{
+                clients: map[string]pubsubClient{},
+        }
+}
+
 // prodPubSubClient is a wrapper around the production pubsub client.
 type prodPubSubClient struct {
-        ctx    context.Context
-        client *pubsub.Client
+        *pubsub.Client
 }
 
-func (pc *prodPubSubClient) getTopic(id string) (*pubsub.Topic, error) {
-        topic := pc.client.Topic(id)
-        exists, err := topic.Exists(pc.ctx)
+func (pc *prodPubSubClient) getTopic(c context.Context, id string) (*pubsub.Topic, error) {
+        topic := pc.Client.Topic(id)
+        exists, err := topic.Exists(c)
         switch {
         case err != nil:
                 return nil, err
         case !exists:
                 return nil, errNotExist
         }
         return topic, nil
 }
 
-func (pc *prodPubSubClient) getSubscription(id string) (*pubsub.Subscription, error) {
-        sub := pc.client.Subscription(id)
-        exists, err := sub.Exists(pc.ctx)
+func (pc *prodPubSubClient) getSubscription(c context.Context, id string) (*pubsub.Subscription, error) {
+        sub := pc.Client.Subscription(id)
+        exists, err := sub.Exists(c)
         switch {
         case err != nil:
                 return nil, err
         case !exists:
                 return nil, errNotExist
         }
         return sub, nil
 }
 
-func (pc *prodPubSubClient) createSubscription(id string, cfg pubsub.SubscriptionConfig) (
+func (pc *prodPubSubClient) createSubscription(
+        c context.Context, id string, cfg pubsub.SubscriptionConfig) (
         *pubsub.Subscription, error) {
 
-        return pc.client.CreateSubscription(pc.ctx, id, cfg)
+        return pc.Client.CreateSubscription(c, id, cfg)
 }
 
 // getPubSubClient extracts a debug PubSub client out of the context.
-func getPubSubClient(c context.Context) (pubsubClient, error) {
-        if client, ok := c.Value(&pubSubClientKey).(pubsubClient); ok {
-                return client, nil
+func getPubSubClient(c context.Context, projectID string) (pubsubClient, error) {
+        v := c.Value(&pubSubClientKey)
+        if clients, ok := v.(*pubsubClients); ok {
+                clients.lock.RLock()
+                client, ok := clients.clients[projectID]
+                clients.lock.RUnlock()
+                if ok {
+                        return client, nil
+                }
+                return nil, fmt.Errorf("no pubsub clients installed for %s", projectID)
         }
         return nil, errors.New("no pubsub clients installed")
 }
 
 // withClient returns a context with a pubsub client instantiated to the
 // given project ID
 func withClient(c context.Context, projectID string) (context.Context, error) {
         if projectID == "" {
-                return nil, errors.New("missing buildbucket project")
+                return nil, errors.New("missing project id")
         }
         client, err := pubsub.NewClient(c, projectID)

[nodir, 2017/07/13 02:07:11]

Just noticed c here. Memory leakish problem happens here too

if EnsurePubSubSubscribed is the only caller of this function, it might be
better to put the filled map without a mutex into the context right there
without withClient's generality (which seems to be hard to implement correctly)

However, a global pubsub client can be problematic. After reading pubsub code a
bit, I learned that if someone calls
http://godoc.org/cloud.google.com/go/pubsub#Client.Close on a global client, it
will affect all future requests that use the client

Stepping back, I assume clients are in the context because you want to use stubs
in tests. Same result could be achieved by a simpler approach: store a func(c
context.Context, projectID string) pubSubClient
In context.

In getPubSubClient, check if function is there. If it is, use it, otherwise
create a prod instance. Then you don't need pubsubClients, the map, the mutex,
newClientsBundle or withClient. You wouldn't really need getTopic,
getSubscription and createSubscription either. You could call getPubSubClient
once and then use its methods directly. I think overall it would lead to less
code

Up to you

[Ryan Tseng, 2017/07/13 02:27:18]

The client is in the context because of the guideline:
https://godoc.org/cloud.google.com/go/pubsub#Client
"Clients should be reused rather than being created as needed"

I assume that is because creating a client also creates 2 RPC connection dials:
https://github.com/GoogleCloudPlatform/google-cloud-go/blob/master/pubsub/api...
https://github.com/GoogleCloudPlatform/google-cloud-go/blob/master/pubsub/api...

We do reuse the client for multiple calls, so creating multiple connection for
each call wouldn't be super elegant.  I'm not super worried about the global
.Close() if we don't expose the underlying client.  The embedded context is
unfortunate though, since there's no other way to reuse an existing client than
to hold a reference to the context at that time

         if err != nil {
                 return nil, err
         }
-        return context.WithValue(c, &pubSubClientKey, &prodPubSubClient{c, client}), nil
+        var clients *pubsubClients
+        var ok bool
+        if clients, ok = c.Value(&pubSubClientKey).(*pubsubClients); !ok {
+                clients = newClientsBundle()
+                c = context.WithValue(c, &pubSubClientKey, clients)
+        }
+        clients.lock.Lock()
+        clients.clients[projectID] = &prodPubSubClient{client}

[nodir, 2017/07/13 02:07:11]

What if it is already there?

[Ryan Tseng, 2017/07/13 02:27:18]

it's overwritten

+        clients.lock.Unlock()
+        return c, nil
 }
 
-func getTopic(c context.Context, id string) (*pubsub.Topic, error) {
-        client, err := getPubSubClient(c)
+func getTopic(c context.Context, project, id string) (*pubsub.Topic, error) {
+        client, err := getPubSubClient(c, project)
         if err != nil {
                 return nil, err
         }
-        return client.getTopic(id)
+        return client.getTopic(c, id)
 }
 
-func getSubscription(c context.Context, id string) (*pubsub.Subscription, error) {
-        client, err := getPubSubClient(c)
+func getSubscription(c context.Context, project, id string) (*pubsub.Subscription, error) {
+        client, err := getPubSubClient(c, project)
         if err != nil {
                 return nil, err
         }
-        return client.getSubscription(id)
+        return client.getSubscription(c, id)
 }
 
-func createSubscription(c context.Context, id string, cfg pubsub.SubscriptionConfig) (
+func createSubscription(c context.Context, project, id string, cfg pubsub.SubscriptionConfig) (
         *pubsub.Subscription, error) {
 
-        client, err := getPubSubClient(c)
+        client, err := getPubSubClient(c, project)
         if err != nil {
                 return nil, err
         }
-        return client.createSubscription(id, cfg)
+        return client.createSubscription(c, id, cfg)
 }
 
 // EnsurePubSubSubscribed makes sure the following subscriptions are in place:
 // * buildbucket, via the settings.Buildbucket.Topic setting
 func EnsurePubSubSubscribed(c context.Context, settings *config.Settings) error {
         if settings.Buildbucket != nil {
                 // Install the production pubsub client pointing to the buildbucket project
-                // into the context.
+                // into the context, so that we can get a reference to the topic.
                 c, err := withClient(c, settings.Buildbucket.Project)
                 if err != nil {
                         return err
                 }
+                // We also need a client for this project for the subscription.
+                c, err = withClient(c, info.AppID(c))
+                if err != nil {
+                        return err
+                }
                 return ensureBuildbucketSubscribed(c, settings.Buildbucket.Project)
         }
         // TODO(hinoka): Ensure buildbot subscribed.
         return nil
 }
 
 // ensureSubscribed is called by a cron job and ensures that the Milo
 // instance is properly subscribed to the buildbucket subscription endpoint.
 func ensureBuildbucketSubscribed(c context.Context, projectID string) error {
         topicID := "builds"
         // Check to see if the topic exists first.
-        topic, err := getTopic(c, topicID)
+        topic, err := getTopic(c, projectID, topicID)
         switch err {
         case errNotExist:
-                logging.WithError(err).Errorf(c, "%s does not exist", topicID)
-                return err
+                return luciErrors.Annotate(err, "%s does not exist", topicID).Err()
         case nil:
                 // continue
         default:
                 if strings.Contains(err.Error(), "PermissionDenied") {
                         URL := "https://console.cloud.google.com/iam-admin/iam/project?project=" + projectID
                         acct, serr := info.ServiceAccount(c)
                         if serr != nil {
                                 acct = fmt.Sprintf("Unknown: %s", serr.Error())
                         }
                         // The documentation is incorrect.  We need Editor permission because
                         // the Subscriber permission does NOT permit attaching subscriptions to
                         // topics or to view topics.
                         logging.WithError(err).Errorf(
                                 c, "please go to %s and add %s as a Pub/Sub Editor", URL, acct)
                 } else {
                         logging.WithError(err).Errorf(c, "could not check topic %#v", topic)
                 }
                 return err
         }
         // Now check to see if the subscription already exists.
-        subID := info.AppID(c)
-        // Get the pubsub module of our app.  We do not want to use info.ModuleHostname()
-        // because it returns a version pinned hostname instead of the default route.
-        sub, err := getSubscription(c, subID)
+        appID := info.AppID(c)
+        sub, err := getSubscription(c, appID, "buildbucket")
         switch err {
         case errNotExist:
                 // continue
         case nil:
                 logging.Infof(c, "subscription %#v exists, no need to update", sub)
                 return nil
         default:
                 logging.WithError(err).Errorf(c, "could not check subscription %#v", sub)
                 return err
         }
+        // Get the pubsub module of our app.  We do not want to use info.ModuleHostname()
+        // because it returns a version pinned hostname instead of the default route.
         pubsubModuleHost := "pubsub." + info.DefaultVersionHostname(c)
 
         // No subscription exists, attach a new subscription to the existing topic.
         endpointURL := url.URL{
                 Scheme: "https",
                 Host:   pubsubModuleHost,
                 Path:   "/_ah/push-handlers/buildbucket",
         }
         subConfig := pubsub.SubscriptionConfig{
                 Topic:       topic,
                 PushConfig:  pubsub.PushConfig{Endpoint: endpointURL.String()},
                 AckDeadline: time.Minute * 10,
         }
-        newSub, err := createSubscription(c, subID, subConfig)
+        newSub, err := createSubscription(c, appID, "buildbucket", subConfig)
         if err != nil {
-                if strings.Contains(err.Error(), "The supplied HTTP URL is not registered") {
-                        registerURL := "https://console.cloud.google.com/apis/credentials/domainverification?project=" + projectID
-                        verifyURL := "https://www.google.com/webmasters/verification/verification?hl=en-GB&siteUrl=http://" + pubsubModuleHost
-                        logging.WithError(err).Errorf(
-                                c, "The domain has to be verified and added.\n\n"+
-                                        "1. Go to %s\n"+
-                                        "2. Verify the domain\n"+
-                                        "3. Go to %s\n"+
-                                        "4. Add %s to allowed domains\n\n",
-                                verifyURL, registerURL, pubsubModuleHost)
-                } else {
-                        logging.WithError(err).Errorf(c, "could not create subscription %#v", sub)
-                }
-                return err
+                return luciErrors.Annotate(err, "could not create subscription %#v", sub).Err()
         }
         // Success!
         logging.Infof(c, "successfully created subscription %#v", newSub)
         return nil
 }