LogDog: BigTable batching schema.

server/internal/logdog/service/service.go

 // Copyright 2016 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 service
 
 import (
         "errors"
         "flag"
         "net/http"
         "os"
         "os/signal"
+        "runtime/pprof"
         "sync/atomic"
 
         "github.com/luci/luci-go/client/authcli"
         "github.com/luci/luci-go/common/api/logdog_coordinator/services/v1"
         "github.com/luci/luci-go/common/auth"
         "github.com/luci/luci-go/common/gcloud/gs"
         log "github.com/luci/luci-go/common/logging"
         "github.com/luci/luci-go/common/logging/gologger"
         "github.com/luci/luci-go/common/proto/logdog/svcconfig"
         "github.com/luci/luci-go/common/prpc"
@@ skipping 24 matching lines @@
 
         shutdownFunc atomic.Value
 
         loggingFlags log.Config
         authFlags    authcli.Flags
         configFlags  config.Flags
 
         coordinatorHost           string
         coordinatorInsecure       bool
         storageCredentialJSONPath string
+        cpuProfilePath            string
 
         coord  logdog.ServicesClient
         config *config.Manager
 }
 
 // Run performs service-wide initialization and invokes the specified run
 // function.
 func (s *Service) Run(c context.Context, f func(context.Context) error) {
         c = gologger.Use(c)
 
         rc := 0
         if err := s.runImpl(c, f); err != nil {
                 log.WithError(err).Errorf(c, "Application exiting with error.")
                 rc = 1
         }
         os.Exit(rc)
 }
 
 func (s *Service) runImpl(c context.Context, f func(context.Context) error) error {
         s.addFlags(c, &s.Flags)
         if err := s.Flags.Parse(os.Args[1:]); err != nil {
                 log.WithError(err).Errorf(c, "Failed to parse command-line.")
                 return err
         }
 
         // Install logging configuration.
         c = s.loggingFlags.Set(c)
 
+        if p := s.cpuProfilePath; p != "" {
+                fd, err := os.Create(p)
+                if err != nil {
+                        log.Fields{
+                                log.ErrorKey: err,
+                                "path":       p,
+                        }.Errorf(c, "Failed to create CPU profile output file.")
+                        return err
+                }
+                defer fd.Close()
+
+                pprof.StartCPUProfile(fd)
+                defer pprof.StopCPUProfile()
+        }
+
         // Configure our signal handler. It will listen for terminating signals and
         // issue a shutdown signal if one is received.
         signalC := make(chan os.Signal)
         go func() {
                 hasShutdownAlready := false
                 for sig := range signalC {
                         if !hasShutdownAlready {
                                 hasShutdownAlready = true
 
                                 log.Warningf(log.SetField(c, "signal", sig), "Received close signal. Send again to terminate immediately.")
@@ skipping 40 matching lines @@
                 Logger:  log.Get(c),
         })
         s.configFlags.AddToFlagSet(fs)
 
         fs.StringVar(&s.coordinatorHost, "coordinator", "",
                 "The Coordinator service's [host][:port].")
         fs.BoolVar(&s.coordinatorInsecure, "coordinator-insecure", false,
                 "Connect to Coordinator over HTTP (instead of HTTPS).")
         fs.StringVar(&s.storageCredentialJSONPath, "storage-credential-json-path", "",
                 "If supplied, the path of a JSON credential file to load and use for storage operations.")
+        fs.StringVar(&s.cpuProfilePath, "cpu-profile-path", "",
+                "If supplied, enable CPU profiling and write the profile here.")
 }
 
 func (s *Service) initCoordinatorClient(c context.Context) (logdog.ServicesClient, error) {
         if s.coordinatorHost == "" {
                 log.Errorf(c, "Missing Coordinator URL (-coordinator).")
                 return nil, ErrInvalidConfig
         }
 
         httpClient, err := s.AuthenticatedClient(func(o *auth.Options) {
                 o.Scopes = CoordinatorScopes
@@ skipping 83 matching lines @@
                 o.Scopes = bigtable.StorageScopes
                 if s.storageCredentialJSONPath != "" {
                         o.ServiceAccountJSONPath = s.storageCredentialJSONPath
                 }
         })
         if err != nil {
                 log.WithError(err).Errorf(c, "Failed to create BigTable Authenticator.")
                 return nil, err
         }
 
-        return bigtable.New(c, bigtable.Options{
+        bt, err := bigtable.New(c, bigtable.Options{
                 Project:  btcfg.Project,
                 Zone:     btcfg.Zone,
                 Cluster:  btcfg.Cluster,
                 LogTable: btcfg.LogTableName,
                 ClientOptions: []cloud.ClientOption{
                         cloud.WithTokenSource(a.TokenSource()),
                 },
-        }), nil
+        })
+        if err != nil {
+                return nil, err
+        }
+        return bt, nil
 }
 
 // GSClient returns an authenticated Google Storage client instance.
 func (s *Service) GSClient(c context.Context) (gs.Client, error) {
         rt, err := s.AuthenticatedTransport(func(o *auth.Options) {
                 o.Scopes = gs.ReadWriteScopes
         })
         if err != nil {
                 log.WithError(err).Errorf(c, "Failed to create authenticated GS transport.")
                 return nil, err
@@ skipping 41 matching lines @@
 //
 // An optional permutation functon can be provided to modify those Options
 // before the Authenticator is created.
 func (s *Service) AuthenticatedClient(f func(o *auth.Options)) (*http.Client, error) {
         a, err := s.Authenticator(f)
         if err != nil {
                 return nil, err
         }
         return a.Client()
 }