X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/533ee50c604e3900c390b77fafac1455b28a15d0..db118cca358662e57a3dd0c1186dce1f0a62ca52:/lib/controller/handler.go

diff --git a/lib/controller/handler.go b/lib/controller/handler.go
index e1392bef92..5d6bd9570c 100644
--- a/lib/controller/handler.go
+++ b/lib/controller/handler.go
@@ -6,7 +6,6 @@ package controller
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
@@ -21,10 +20,8 @@ import (
 	"git.arvados.org/arvados.git/lib/controller/router"
 	"git.arvados.org/arvados.git/lib/ctrlctx"
 	"git.arvados.org/arvados.git/sdk/go/arvados"
-	"git.arvados.org/arvados.git/sdk/go/ctxlog"
 	"git.arvados.org/arvados.git/sdk/go/health"
 	"git.arvados.org/arvados.git/sdk/go/httpserver"
-	"github.com/jmoiron/sqlx"
 
 	// sqlx needs lib/pq to talk to PostgreSQL
 	_ "github.com/lib/pq"
@@ -40,8 +37,8 @@ type Handler struct {
 	proxy          *proxy
 	secureClient   *http.Client
 	insecureClient *http.Client
-	pgdb           *sqlx.DB
-	pgdbMtx        sync.Mutex
+	dbConnector    ctrlctx.DBConnector
+	limitLogCreate chan struct{}
 }
 
 func (h *Handler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
@@ -65,7 +62,7 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 
 func (h *Handler) CheckHealth() error {
 	h.setupOnce.Do(h.setup)
-	_, err := h.db(context.TODO())
+	_, err := h.dbConnector.GetDB(context.TODO())
 	if err != nil {
 		return err
 	}
@@ -97,17 +94,18 @@ func (h *Handler) setup() {
 	mux := http.NewServeMux()
 	healthFuncs := make(map[string]health.Func)
 
-	oidcAuthorizer := localdb.OIDCAccessTokenAuthorizer(h.Cluster, h.db)
+	h.dbConnector = ctrlctx.DBConnector{PostgreSQL: h.Cluster.PostgreSQL}
+	oidcAuthorizer := localdb.OIDCAccessTokenAuthorizer(h.Cluster, h.dbConnector.GetDB)
 	h.federation = federation.New(h.Cluster, &healthFuncs)
 	rtr := router.New(h.federation, router.Config{
 		MaxRequestSize: h.Cluster.API.MaxRequestSize,
 		WrapCalls: api.ComposeWrappers(
-			ctrlctx.WrapCallsInTransactions(h.db),
+			ctrlctx.WrapCallsInTransactions(h.dbConnector.GetDB),
 			oidcAuthorizer.WrapCalls,
 			ctrlctx.WrapCallsWithAuth(h.Cluster)),
 	})
 
-	healthRoutes := health.Routes{"ping": func() error { _, err := h.db(context.TODO()); return err }}
+	healthRoutes := health.Routes{"ping": func() error { _, err := h.dbConnector.GetDB(context.TODO()); return err }}
 	for name, f := range healthFuncs {
 		healthRoutes[name] = f
 	}
@@ -137,6 +135,7 @@ func (h *Handler) setup() {
 
 	hs := http.NotFoundHandler()
 	hs = prepend(hs, h.proxyRailsAPI)
+	hs = prepend(hs, h.limitLogCreateRequests)
 	hs = h.setupProxyRemoteCluster(hs)
 	hs = prepend(hs, oidcAuthorizer.Middleware)
 	mux.Handle("/", hs)
@@ -150,6 +149,12 @@ func (h *Handler) setup() {
 	ic.CheckRedirect = neverRedirect
 	h.insecureClient = &ic
 
+	logCreateLimit := int(float64(h.Cluster.API.MaxConcurrentRequests) * h.Cluster.API.LogCreateRequestFraction)
+	if logCreateLimit == 0 && h.Cluster.API.LogCreateRequestFraction > 0 {
+		logCreateLimit = 1
+	}
+	h.limitLogCreate = make(chan struct{}, logCreateLimit)
+
 	h.proxy = &proxy{
 		Name: "arvados-controller",
 	}
@@ -158,31 +163,6 @@ func (h *Handler) setup() {
 	go h.containerLogSweepWorker()
 }
 
-var errDBConnection = errors.New("database connection error")
-
-func (h *Handler) db(ctx context.Context) (*sqlx.DB, error) {
-	h.pgdbMtx.Lock()
-	defer h.pgdbMtx.Unlock()
-	if h.pgdb != nil {
-		return h.pgdb, nil
-	}
-
-	db, err := sqlx.Open("postgres", h.Cluster.PostgreSQL.Connection.String())
-	if err != nil {
-		ctxlog.FromContext(ctx).WithError(err).Error("postgresql connect failed")
-		return nil, errDBConnection
-	}
-	if p := h.Cluster.PostgreSQL.ConnectionPool; p > 0 {
-		db.SetMaxOpenConns(p)
-	}
-	if err := db.Ping(); err != nil {
-		ctxlog.FromContext(ctx).WithError(err).Error("postgresql connect succeeded but ping failed")
-		return nil, errDBConnection
-	}
-	h.pgdb = db
-	return db, nil
-}
-
 type middlewareFunc func(http.ResponseWriter, *http.Request, http.Handler)
 
 func prepend(next http.Handler, middleware middlewareFunc) http.Handler {
@@ -210,6 +190,20 @@ func (h *Handler) localClusterRequest(req *http.Request) (*http.Response, error)
 	return h.proxy.Do(req, urlOut, client)
 }
 
+func (h *Handler) limitLogCreateRequests(w http.ResponseWriter, req *http.Request, next http.Handler) {
+	if cap(h.limitLogCreate) > 0 && req.Method == http.MethodPost && strings.HasPrefix(req.URL.Path, "/arvados/v1/logs") {
+		select {
+		case h.limitLogCreate <- struct{}{}:
+			defer func() { <-h.limitLogCreate }()
+			next.ServeHTTP(w, req)
+		default:
+			http.Error(w, "Excess log messages", http.StatusServiceUnavailable)
+		}
+		return
+	}
+	next.ServeHTTP(w, req)
+}
+
 func (h *Handler) proxyRailsAPI(w http.ResponseWriter, req *http.Request, next http.Handler) {
 	resp, err := h.localClusterRequest(req)
 	n, err := h.proxy.ForwardResponse(w, resp, err)