LogDog: Re-use Pub/Sub gRPC clients.

logdog/appengine/coordinator/service.go

 // Copyright 2015 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 coordinator
 
 import (
-        "net/http"
         "sync"
         "sync/atomic"
         "time"
 
         "github.com/luci/luci-go/appengine/gaeauth/server/gaesigner"
         "github.com/luci/luci-go/common/clock"
         luciConfig "github.com/luci/luci-go/common/config"
         "github.com/luci/luci-go/common/errors"
         "github.com/luci/luci-go/common/gcloud/gs"
         "github.com/luci/luci-go/common/gcloud/pubsub"
@@ skipping 64 matching lines @@
         StorageForStream(context.Context, *LogStreamState) (Storage, error)
 
         // ArchivalPublisher returns an ArchivalPublisher instance.
         ArchivalPublisher(context.Context) (ArchivalPublisher, error)
 }
 
 // ProdCoordinatorService is Middleware used by Coordinator services.
 //
 // It installs a production Services instance into the Context.
 func ProdCoordinatorService(c *router.Context, next router.Handler) {
-        c.Context = WithServices(c.Context, &prodServicesInst{
-                req: c.Request,
-        })
+        services := prodServicesInst{
+                aeCtx: appengine.WithContext(c.Context, c.Request),
+        }
+        defer services.close(c.Context)
+
+        c.Context = WithServices(c.Context, &services)
         next(c)
 }
 
 // prodServicesInst is a Service exposing production faciliites. A unique
 // instance is bound to each each request.
 type prodServicesInst struct {
         sync.Mutex
 
-        // req is the request that is being serviced. This is populated when the
-        // service is installed via ProdCoordinatorService.
-        req *http.Request
+        // aeCtx is an AppEngine Context initialized for the current request.
+        aeCtx context.Context
 
         // gcfg is the cached global configuration.
         gcfg           *config.Config
         projectConfigs map[luciConfig.ProjectName]*cachedProjectConfig
 
         // archivalIndex is the atomically-manipulated archival index for the
         // ArchivalPublisher. This is shared between all ArchivalPublisher instances
         // from this service.
         archivalIndex int32
 
+        // pubSubClients is a map of Pub/Sub client singletons generated during this
+        // request. Each client is associated with its project, and will be
+        // initialized the first time it is requested by getPubSubClient.
+        //
+        // All clients will be closed on "close".
+        pubSubClients map[string]*gcps.Client
+
         // signer is the signer instance to use.
         signer gaesigner.Signer
 }
 
+func (s *prodServicesInst) close(c context.Context) {
+        s.Lock()
+        defer s.Unlock()
+
+        for proj, client := range s.pubSubClients {
+                if err := client.Close(); err != nil {
+                        log.Fields{
+                                log.ErrorKey: err,
+                                "project":    proj,
+                        }.Errorf(c, "Failed to close Pub/Sub client singleton.")
+                }
+        }
+}
+
 func (s *prodServicesInst) Config(c context.Context) (*config.Config, error) {
         s.Lock()
         defer s.Unlock()
 
         // Load/cache the global config.
         if s.gcfg == nil {
                 var err error
                 s.gcfg, err = config.Load(c)
                 if err != nil {
                         return nil, err
@@ skipping 195 matching lines @@
         fullTopic := pubsub.Topic(cfg.Coordinator.ArchiveTopic)
         if err := fullTopic.Validate(); err != nil {
                 log.Fields{
                         log.ErrorKey: err,
                         "topic":      fullTopic,
                 }.Errorf(c, "Failed to validate archival topic.")
                 return nil, errors.New("invalid archival topic")
         }
         project, topic := fullTopic.Split()
 
+        // Get a Pub/Sub client (maybe re-use).
+        psClient, err := s.getPubSubClient(project)
+        if err != nil {
+                log.WithError(err).Errorf(c, "Failed to get Pub/Sub client.")
+                return nil, err
+        }
+
+        return &pubsubArchivalPublisher{
+                client:           psClient,
+                topic:            psClient.Topic(topic),
+                publishIndexFunc: s.nextArchiveIndex,
+        }, nil
+}
+
+func (s *prodServicesInst) getPubSubClient(proj string) (*gcps.Client, error) {
+        // Create a new AppEngine context. Don't pass gRPC metadata to PubSub, since
+        // we don't want any caller RPC to be forwarded to the backend service.
+        c := metadata.NewContext(s.aeCtx, nil)
+
+        s.Lock()
+        defer s.Unlock()
+
+        client := s.pubSubClients[proj]
+        if client != nil {
+                return client, nil
+        }
+
         // Create an authenticated Pub/Sub client.
         creds, err := auth.GetPerRPCCredentials(auth.AsSelf, auth.WithScopes(pubsub.PublisherScopes...))
         if err != nil {
                 log.WithError(err).Errorf(c, "Failed to create Pub/Sub credentials.")
                 return nil, errors.New("failed to create Pub/Sub credentials")
         }
 
-        // Create a new AppEngine context. Don't pass gRPC metadata to PubSub, since
-        // we don't want any caller RPC to be forwarded to the backend service.
-        aeCtx := appengine.WithContext(metadata.NewContext(c, nil), s.req)
-
-        psClient, err := gcps.NewClient(aeCtx, project,
+        psClient, err := gcps.NewClient(c, proj,
                 option.WithGRPCDialOption(grpc.WithPerRPCCredentials(creds)))
         if err != nil {
                 log.WithError(err).Errorf(c, "Failed to create Pub/Sub client.")
                 return nil, errors.New("failed to create Pub/Sub client")
         }
 
-        return &pubsubArchivalPublisher{
-                client:           psClient,
-                topic:            psClient.Topic(topic),
-                publishIndexFunc: s.nextArchiveIndex,
-        }, nil
+        // Retain this client for the duration of our request.
+        if s.pubSubClients == nil {
+                s.pubSubClients = make(map[string]*gcps.Client)
+        }
+        s.pubSubClients[proj] = psClient
+        return psClient, nil
 }
 
 func (s *prodServicesInst) nextArchiveIndex() uint64 {
         // We use a 32-bit value for this because it avoids atomic memory bounary
         // issues. Furthermore, we constrain it to be positive, using a negative
         // value as a sentinel that the archival index has wrapped.
         //
         // This is reasonable, as it is very unlikely that a single request will issue
         // more than MaxInt32 archival tasks.
         v := atomic.AddInt32(&s.archivalIndex, 1) - 1
@@ skipping 133 matching lines @@
 
         // Sign index URL.
         if req.Index {
                 if resp.Index, err = doSign(si.index); err != nil {
                         return nil, errors.Annotate(err).InternalReason("failed to sign index URL").Err()
                 }
         }
 
         return &resp, nil
 }