Add debug endpoints on export

cmd/export/main.go

@@ -18,7 +18,6 @@ package main
 import (
  "context"
  "fmt"
- "net/http"
 
  "github.com/google/exposure-notifications-server/internal/export"
  "github.com/google/exposure-notifications-server/internal/interrupt"
@@ -53,15 +52,11 @@ func realMain(ctx context.Context) error {
  return fmt.Errorf("export.NewServer: %w", err)
  }
 
- mux := http.NewServeMux()
- mux.HandleFunc("/create-batches", batchServer.CreateBatchesHandler)
- mux.HandleFunc("/do-work", batchServer.WorkerHandler)
-
  srv, err := server.New(config.Port)
  if err != nil {
  return fmt.Errorf("server.New: %w", err)
  }
  logger.Infof("listening on :%s", config.Port)
 
- return srv.ServeHTTPHandler(ctx, mux)
+ return srv.ServeHTTPHandler(ctx, batchServer.Routes(ctx))
 }

internal/database/config.go

@@ -19,21 +19,21 @@ import (
 )
 
 type Config struct {
- Name string `env:"DB_NAME"`
- User string `env:"DB_USER"`
- Host string `env:"DB_HOST, default=localhost"`
- Port string `env:"DB_PORT, default=5432"`
- SSLMode string `env:"DB_SSLMODE, default=require"`
- ConnectionTimeout int `env:"DB_CONNECT_TIMEOUT"`
+ Name string `env:"DB_NAME" json:",omitempty"`
+ User string `env:"DB_USER" json:",omitempty"`
+ Host string `env:"DB_HOST, default=localhost" json:",omitempty"`
+ Port string `env:"DB_PORT, default=5432" json:",omitempty"`
+ SSLMode string `env:"DB_SSLMODE, default=require" json:",omitempty"`
+ ConnectionTimeout int `env:"DB_CONNECT_TIMEOUT" json:",omitempty"`
  Password string `env:"DB_PASSWORD" json:"-"` // ignored by zap's JSON formatter
- SSLCertPath string `env:"DB_SSLCERT"`
- SSLKeyPath string `env:"DB_SSLKEY"`
- SSLRootCertPath string `env:"DB_SSLROOTCERT"`
- PoolMinConnections string `env:"DB_POOL_MIN_CONNS"`
- PoolMaxConnections string `env:"DB_POOL_MAX_CONNS"`
- PoolMaxConnLife time.Duration `env:"DB_POOL_MAX_CONN_LIFETIME"`
- PoolMaxConnIdle time.Duration `env:"DB_POOL_MAX_CONN_IDLE_TIME"`
- PoolHealthCheck time.Duration `env:"DB_POOL_HEALTH_CHECK_PERIOD"`
+ SSLCertPath string `env:"DB_SSLCERT" json:",omitempty"`
+ SSLKeyPath string `env:"DB_SSLKEY" json:",omitempty"`
+ SSLRootCertPath string `env:"DB_SSLROOTCERT" json:",omitempty"`
+ PoolMinConnections string `env:"DB_POOL_MIN_CONNS" json:",omitempty"`
+ PoolMaxConnections string `env:"DB_POOL_MAX_CONNS" json:",omitempty"`
+ PoolMaxConnLife time.Duration `env:"DB_POOL_MAX_CONN_LIFETIME" json:",omitempty"`
+ PoolMaxConnIdle time.Duration `env:"DB_POOL_MAX_CONN_IDLE_TIME" json:",omitempty"`
+ PoolHealthCheck time.Duration `env:"DB_POOL_HEALTH_CHECK_PERIOD" json:",omitempty"`
 }
 
 func (c *Config) DatabaseConfig() *Config {

internal/export/batcher.go

@@ -22,24 +22,27 @@ import (
  "time"
 
  coredb "github.com/google/exposure-notifications-server/internal/database"
- "github.com/google/exposure-notifications-server/internal/export/database"
+ exportdatabase "github.com/google/exposure-notifications-server/internal/export/database"
  "github.com/google/exposure-notifications-server/internal/export/model"
  publishmodel "github.com/google/exposure-notifications-server/internal/publish/model"
 
  "github.com/google/exposure-notifications-server/internal/logging"
 )
 
-// CreateBatchesHandler is a handler to iterate the rows of ExportConfig and
+// handleCreateBatches is a handler to iterate the rows of ExportConfig and
 // create entries in ExportBatchJob as appropriate.
-func (s *Server) CreateBatchesHandler(w http.ResponseWriter, r *http.Request) {
- ctx, cancel := context.WithTimeout(r.Context(), s.config.CreateTimeout)
- defer cancel()
+func (s *Server) handleCreateBatches(ctx context.Context) http.HandlerFunc {
  logger := logging.FromContext(ctx)
  metrics := s.env.MetricsExporter(ctx)
+ db := s.env.Database()
+
+ return func(w http.ResponseWriter, r *http.Request) {
+ ctx, cancel := context.WithTimeout(r.Context(), s.config.CreateTimeout)
+ defer cancel()
 
  // Obtain lock to make sure there are no other processes working to create batches.
  lock := "create_batches"
- unlockFn, err := s.db.Lock(ctx, lock, s.config.CreateTimeout)
+  unlockFn, err := db.Lock(ctx, lock, s.config.CreateTimeout)
  if err != nil {
  if errors.Is(err, coredb.ErrAlreadyLocked) {
  metrics.WriteInt("export-batcher-lock-contention", true, 1)
@@ -66,7 +69,7 @@ func (s *Server) CreateBatchesHandler(w http.ResponseWriter, r *http.Request) {
  }()
 
  effectiveTime := time.Now().Add(-1 * s.config.MinWindowAge)
- err = database.New(s.db).IterateExportConfigs(ctx, effectiveTime, func(ec *model.ExportConfig) error {
+  err = exportdatabase.New(db).IterateExportConfigs(ctx, effectiveTime, func(ec *model.ExportConfig) error {
  totalConfigs++
  if batchesCreated, err := s.maybeCreateBatches(ctx, ec, effectiveTime); err != nil {
  logger.Errorf("Failed to create batches for config %d: %v, continuing to next config", ec.ConfigID, err)
@@ -94,12 +97,14 @@ func (s *Server) CreateBatchesHandler(w http.ResponseWriter, r *http.Request) {
  http.Error(w, "Failed to create batches, check logs.", http.StatusInternalServerError)
  }
  }
+}
 
 func (s *Server) maybeCreateBatches(ctx context.Context, ec *model.ExportConfig, now time.Time) (int, error) {
  logger := logging.FromContext(ctx)
  metrics := s.env.MetricsExporter(ctx)
+ db := s.env.Database()
 
- latestEnd, err := database.New(s.db).LatestExportBatchEnd(ctx, ec)
+ latestEnd, err := exportdatabase.New(db).LatestExportBatchEnd(ctx, ec)
  if err != nil {
  return 0, fmt.Errorf("fetching most recent batch for config %d: %w", ec.ConfigID, err)
  }
@@ -128,7 +133,7 @@ func (s *Server) maybeCreateBatches(ctx context.Context, ec *model.ExportConfig,
  })
  }
 
- if err := database.New(s.db).AddExportBatches(ctx, batches); err != nil {
+ if err := exportdatabase.New(db).AddExportBatches(ctx, batches); err != nil {
  return 0, fmt.Errorf("creating export batches for config %d: %w", ec.ConfigID, err)
  }
 

internal/export/config.go

@@ -48,6 +48,12 @@ type Config struct {
  TruncateWindow time.Duration `env:"TRUNCATE_WINDOW, default=1h"`
  MinWindowAge time.Duration `env:"MIN_WINDOW_AGE, default=2h"`
  TTL time.Duration `env:"CLEANUP_TTL, default=336h"`
+
+ // Debugging flags follow. These should not be enabled in production
+ // environments.
+
+ // DebugEndpointEnabled controls whether the debug endpoint is enabled.
+ DebugEndpoint bool `env:"DEBUG_ENDPOINT"`
 }
 
 func (c *Config) BlobstoreConfig() *storage.Config {

internal/export/database/export.go

@@ -125,7 +125,7 @@ func (db *ExportDB) GetAllExportConfigs(ctx context.Context) ([]*model.ExportCon
  }
  defer rows.Close()
 
- results := []*model.ExportConfig{}
+ var results []*model.ExportConfig
  for rows.Next() {
  ec, err := scanOneExportConfig(rows)
  if err != nil {

internal/export/debug.go

@@ -0,0 +1,97 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package export
+
+import (
+ "context"
+ "encoding/json"
+ "fmt"
+ "net/http"
+ "time"
+
+ exportdatabase "github.com/google/exposure-notifications-server/internal/export/database"
+ exportmodel "github.com/google/exposure-notifications-server/internal/export/model"
+ "github.com/google/exposure-notifications-server/internal/logging"
+)
+
+func (s *Server) handleDebug(ctx context.Context) http.HandlerFunc {
+ logger := logging.FromContext(ctx)
+
+ type response struct {
+ Config *Config
+ ExportConfigs []*exportmodel.ExportConfig
+ ExportBatchEnds map[int64]time.Time
+ ExportBatchFiles []string
+
+ SignatureInfos []*exportmodel.SignatureInfo
+ }
+
+ return func(w http.ResponseWriter, r *http.Request) {
+ ctx := r.Context()
+ db := s.env.Database()
+
+ exportConfigs, err := exportdatabase.New(db).GetAllExportConfigs(ctx)
+ if err != nil {
+ logger.Errorf("failed to get all export configs: %v", err)
+ w.WriteHeader(http.StatusInternalServerError)
+ fmt.Fprint(w, http.StatusText(http.StatusInternalServerError))
+ return
+ }
+
+ exportBatchEnds := make(map[int64]time.Time, len(exportConfigs))
+ for _, ec := range exportConfigs {
+ end, err := exportdatabase.New(db).LatestExportBatchEnd(ctx, ec)
+ if err != nil {
+ logger.Errorf("failed to get latest export batch end for %d: %v", ec.ConfigID, err)
+ w.WriteHeader(http.StatusInternalServerError)
+ fmt.Fprint(w, http.StatusText(http.StatusInternalServerError))
+ return
+ }
+ exportBatchEnds[ec.ConfigID] = end
+ }
+
+ exportBatchFiles, err := exportdatabase.New(db).LookupExportFiles(ctx, 4*time.Hour)
+ if err != nil {
+ logger.Errorf("failed to get export files: %v", err)
+ w.WriteHeader(http.StatusInternalServerError)
+ fmt.Fprint(w, http.StatusText(http.StatusInternalServerError))
+ return
+ }
+
+ signatureInfos, err := exportdatabase.New(db).ListAllSigntureInfos(ctx)
+ if err != nil {
+ logger.Errorf("failed to get all signature infos: %v", err)
+ w.WriteHeader(http.StatusInternalServerError)
+ fmt.Fprint(w, http.StatusText(http.StatusInternalServerError))
+ return
+ }
+
+ resp := &response{
+ Config: s.config,
+ ExportConfigs: exportConfigs,
+ ExportBatchEnds: exportBatchEnds,
+ ExportBatchFiles: exportBatchFiles,
+ SignatureInfos: signatureInfos,
+ }
+
+ w.Header().Set("Content-Type", "application/json")
+
+ e := json.NewEncoder(w)
+ e.SetIndent("", " ")
+ if err := e.Encode(resp); err != nil {
+ panic(err)
+ }
+ }
+}

internal/export/server.go

@@ -15,11 +15,9 @@
 package export
 
 import (
+ "context"
  "fmt"
-
- coredb "github.com/google/exposure-notifications-server/internal/database"
- "github.com/google/exposure-notifications-server/internal/export/database"
- publishdb "github.com/google/exposure-notifications-server/internal/publish/database"
+ "net/http"
 
  "github.com/google/exposure-notifications-server/internal/serverenv"
 )
@@ -30,6 +28,9 @@ func NewServer(config *Config, env *serverenv.ServerEnv) (*Server, error) {
  if env.Blobstore() == nil {
  return nil, fmt.Errorf("export.NewBatchServer requires Blobstore present in the ServerEnv")
  }
+ if env.Database() == nil {
+ return nil, fmt.Errorf("export.NewBatchServer requires Database present in the ServerEnv")
+ }
  if env.KeyManager() == nil {
  return nil, fmt.Errorf("export.NewBatchServer requires KeyManager present in the ServerEnv")
  }
@@ -38,19 +39,28 @@ func NewServer(config *Config, env *serverenv.ServerEnv) (*Server, error) {
  }
 
  return &Server{
- db: env.Database(),
- exportdb: database.New(env.Database()),
- publishdb: publishdb.New(env.Database()),
  config: config,
  env: env,
  }, nil
 }
 
 // Server hosts end points to manage export batches.
 type Server struct {
- db *coredb.DB
- exportdb *database.ExportDB
- publishdb *publishdb.PublishDB
  config *Config
  env *serverenv.ServerEnv
 }
+
+// Routes defines and returns the routes for this server.
+func (s *Server) Routes(ctx context.Context) *http.ServeMux {
+ mux := http.NewServeMux()
+
+ mux.HandleFunc("/create-batches", s.handleCreateBatches(ctx))
+ mux.HandleFunc("/do-work", s.handleDoWork(ctx))
+
+ // Enable debug endpoint if configured.
+ if s.config.DebugEndpoint {
+ mux.HandleFunc("/debug", s.handleDebug(ctx))
+ }
+
+ return mux
+}

internal/export/server_test.go

@@ -42,14 +42,28 @@ func TestNewServer(t *testing.T) {
  env: serverenv.New(ctx),
  err: fmt.Errorf("export.NewBatchServer requires Blobstore present in the ServerEnv"),
  },
+ {
+ name: "nil Database",
+ env: serverenv.New(ctx,
+ serverenv.WithBlobStorage(emptyStorage),
+ ),
+ err: fmt.Errorf("export.NewBatchServer requires Database present in the ServerEnv"),
+ },
  {
  name: "nil KeyManager",
- env: serverenv.New(ctx, serverenv.WithBlobStorage(emptyStorage)),
+ env: serverenv.New(ctx,
+ serverenv.WithBlobStorage(emptyStorage),
+ serverenv.WithDatabase(emptyDB),
+ ),
  err: fmt.Errorf("export.NewBatchServer requires KeyManager present in the ServerEnv"),
  },
  {
  name: "Fully Specified",
- env: serverenv.New(ctx, serverenv.WithBlobStorage(emptyStorage), serverenv.WithKeyManager(emptyKMS), serverenv.WithDatabase(emptyDB)),
+ env: serverenv.New(ctx,
+ serverenv.WithBlobStorage(emptyStorage),
+ serverenv.WithDatabase(emptyDB),
+ serverenv.WithKeyManager(emptyKMS),
+ ),
  err: nil,
  },
  }