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

milo/common/pubsub.go

Index: milo/common/pubsub.go
diff --git a/milo/common/pubsub.go b/milo/common/pubsub.go
index 56b79c2b05902c61395af9853b59294f0865550e..d7f6f73b7cf443f15622526135229d8e09974d71 100644
--- a/milo/common/pubsub.go
+++ b/milo/common/pubsub.go
@@ -6,12 +6,14 @@ import (
 	"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"
 )
@@ -42,25 +44,35 @@ func (m *PubSubSubscription) GetData() ([]byte, error) {
 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
@@ -70,9 +82,9 @@ func (pc *prodPubSubClient) getTopic(id string) (*pubsub.Topic, error) {
 	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
@@ -82,16 +94,24 @@ func (pc *prodPubSubClient) getSubscription(id string) (*pubsub.Subscription, er
 	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")
 }
@@ -100,39 +120,48 @@ func getPubSubClient(c context.Context) (pubsubClient, error) {
 // 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:
@@ -140,11 +169,16 @@ func createSubscription(c context.Context, id string, cfg pubsub.SubscriptionCon
 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.
@@ -156,11 +190,10 @@ func EnsurePubSubSubscribed(c context.Context, settings *config.Settings) error
 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:
@@ -181,10 +214,8 @@ func ensureBuildbucketSubscribed(c context.Context, projectID string) error {
 		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
@@ -195,6 +226,8 @@ func ensureBuildbucketSubscribed(c context.Context, projectID string) error {
 		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.
@@ -208,22 +241,9 @@ func ensureBuildbucketSubscribed(c context.Context, projectID string) error {
 		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)