internal/frontend: recycle database connections every 5m

In order to avoid imbalance between pkgsite's two database instances,
recycle connections every 5 minutes.

Change-Id: I9ca1e686a90f8c61619fd76454ec66163e501ee1
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/680175
kokoro-CI: kokoro <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Jonathan Amsterdam <[email protected]>

Author: findleyr

cmd/frontend/main.go

@@ -33,6 +33,7 @@ import (
 	"golang.org/x/pkgsite/internal/proxy"
 	"golang.org/x/pkgsite/internal/queue"
 	"golang.org/x/pkgsite/internal/queue/gcpqueue"
+	"golang.org/x/pkgsite/internal/resource"
 	"golang.org/x/pkgsite/internal/source"
 	"golang.org/x/pkgsite/internal/static"
 	"golang.org/x/pkgsite/internal/trace"
@@ -70,7 +71,7 @@ func main() {
 	}
 
 	var (
-		dsg        func(context.Context) internal.DataSource
+		dsg        func(context.Context) (internal.DataSource, func())
 		fetchQueue queue.Queue
 	)
 	if *bypassLicenseCheck {
@@ -96,14 +97,19 @@ func main() {
 			ProxyClientForLatest: proxyClient,
 			BypassLicenseCheck:   *bypassLicenseCheck,
 		}.New()
-		dsg = func(context.Context) internal.DataSource { return ds }
+		dsg = func(context.Context) (internal.DataSource, func()) { return ds, func() {} }
 	} else {
-		db, err := cmdconfig.OpenDB(ctx, cfg, *bypassLicenseCheck)
-		if err != nil {
-			log.Fatalf(ctx, "%v", err)
+		dbResource := resource.New(func() *postgres.DB {
+			db, err := cmdconfig.OpenDB(ctx, cfg, *bypassLicenseCheck)
+			if err != nil {
+				log.Fatalf(ctx, "%v", err)
+			}
+			return db
+		}, 5*time.Minute)
+
+		dsg = func(ctx context.Context) (internal.DataSource, func()) {
+			return dbResource.Get()
 		}
-		defer db.Close()
-		dsg = func(context.Context) internal.DataSource { return db }
 		sourceClient := source.NewClient(&http.Client{
 			Transport: new(ochttp.Transport),
 			Timeout:   config.SourceTimeout,
@@ -113,6 +119,8 @@ func main() {
 		// per-request connection.
 		fetchQueue, err = gcpqueue.New(ctx, cfg, queueName, *workers, expg,
 			func(ctx context.Context, modulePath, version string) (int, error) {
+				db, release := dbResource.Get()
+				defer release()
 				return fetchserver.FetchAndUpdateState(ctx, modulePath, version, proxyClient, sourceClient, db)
 			})
 		if err != nil {

cmd/internal/pkgsite/server.go

@@ -47,6 +47,8 @@ type ServerConfig struct {
 }
 
 // BuildServer builds a *frontend.Server using the given configuration.
+//
+// TODO(rfindley): move to the cmd/pkgsite package, its only caller.
 func BuildServer(ctx context.Context, serverCfg ServerConfig) (*frontend.Server, error) {
 	if len(serverCfg.Paths) == 0 && !serverCfg.UseCache && serverCfg.Proxy == nil {
 		serverCfg.Paths = []string{"."}
@@ -288,7 +290,7 @@ func newServer(getters []fetch.ModuleGetter, localModules []frontend.LocalModule
 	go lds.GetUnitMeta(context.Background(), "", "std", "latest")
 
 	server, err := frontend.NewServer(frontend.ServerConfig{
-		DataSourceGetter: func(context.Context) internal.DataSource { return lds },
+		DataSourceGetter: func(context.Context) (internal.DataSource, func()) { return lds, func() {} },
 		TemplateFS:       template.TrustedFSFromEmbed(static.FS),
 		StaticFS:         staticFS,
 		DevMode:          devMode,

internal/frontend/fetchserver/fetch_test.go

@@ -62,7 +62,7 @@ func newTestServerWithFetch(t *testing.T, proxyModules []*proxytest.Module, cach
 
 	s, err := frontend.NewServer(frontend.ServerConfig{
 		FetchServer:      f,
-		DataSourceGetter: func(context.Context) internal.DataSource { return testDB },
+		DataSourceGetter: func(context.Context) (internal.DataSource, func()) { return testDB, func() {} },
 		Queue:            q,
 		TemplateFS:       template.TrustedFSFromEmbed(static.FS),
 		// Use the embedded FSs here to make sure they're tested.

internal/frontend/frontend_test.go

@@ -39,7 +39,7 @@ func newTestServer(t *testing.T, cacher Cacher) (*Server, http.Handler) {
 	t.Helper()
 
 	s, err := NewServer(ServerConfig{
-		DataSourceGetter: func(context.Context) internal.DataSource { return fakedatasource.New() },
+		DataSourceGetter: func(context.Context) (internal.DataSource, func()) { return fakedatasource.New(), func() {} },
 		TemplateFS:       template.TrustedFSFromEmbed(static.FS),
 		// Use the embedded FSs here to make sure they're tested.
 		// Integration tests will use the actual directories.

internal/frontend/latest_version.go

@@ -26,7 +26,8 @@ func (s *Server) GetLatestInfo(ctx context.Context, unitPath, modulePath string,
 
 	// It is okay to use a different DataSource (DB connection) than the rest of the
 	// request, because this makes self-contained calls on the DB.
-	ds := s.getDataSource(ctx)
+	ds, release := s.getDataSource(ctx)
+	defer release()
 
 	latest, err := ds.GetLatestInfo(ctx, unitPath, modulePath, latestUnitMeta)
 	if err != nil {

internal/frontend/latest_version_test.go

@@ -59,7 +59,7 @@ func TestLatestMinorVersion(t *testing.T) {
 	}
 	ctx := context.Background()
 	insertTestModules(ctx, t, fds, persistedModules)
-	svr := &Server{getDataSource: func(context.Context) internal.DataSource { return fds }}
+	svr := &Server{getDataSource: func(context.Context) (internal.DataSource, func()) { return fds, func() {} }}
 	for _, tc := range tt {
 		t.Run(tc.name, func(t *testing.T) {
 			got := svr.GetLatestInfo(ctx, tc.fullPath, tc.modulePath, nil)

internal/frontend/server.go

@@ -42,7 +42,7 @@ import (
 type Server struct {
 	fetchServer FetchServerInterface
 	// getDataSource should never be called from a handler. It is called only in Server.errorHandler.
-	getDataSource      func(context.Context) internal.DataSource
+	getDataSource      func(context.Context) (internal.DataSource, func())
 	queue              queue.Queue
 	templateFS         template.TrustedFS
 	staticFS           fs.FS
@@ -82,9 +82,9 @@ type ServerConfig struct {
 	Config *config.Config
 	// Note that FetchServer may be nil.
 	FetchServer FetchServerInterface
-	// DataSourceGetter should return a DataSource on each call.
+	// DataSourceGetter should return a DataSource and a release function on each call.
 	// It should be goroutine-safe.
-	DataSourceGetter  func(context.Context) internal.DataSource
+	DataSourceGetter  func(context.Context) (internal.DataSource, func())
 	Queue             queue.Queue
 	TemplateFS        template.TrustedFS // for loading templates safely
 	StaticFS          fs.FS              // for static/ directory
@@ -503,7 +503,8 @@ func (s *Server) PanicHandler() (_ http.HandlerFunc, err error) {
 func (s *Server) errorHandler(f func(w http.ResponseWriter, r *http.Request, ds internal.DataSource) error) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		// Obtain a DataSource to use for this request.
-		ds := s.getDataSource(r.Context())
+		ds, release := s.getDataSource(r.Context())
+		defer release()
 		if err := f(w, r, ds); err != nil {
 			s.serveError(w, r, err)
 		}

internal/resource/resource.go

@@ -0,0 +1,92 @@
+// Copyright 2025 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// The resource package defines types to track the lifecycle of transient
+// resources, such as a database connection, that should be renewed at some
+// fixed interval.
+package resource
+
+import (
+	"io"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+type instance[T io.Closer] struct {
+	refs atomic.Int64
+	v    T
+}
+
+func (i *instance[T]) acquire(initial bool) (T, func()) {
+	if i.refs.Add(1) <= 1 && !initial {
+		panic("acquire on released instance")
+	}
+	return i.v, func() {
+		if i.refs.Add(-1) == 0 {
+			i.v.Close() // ignore error
+			var zero T
+			i.v = zero // aid GC
+		}
+	}
+}
+
+// A Resource is a container for a transient resource of type T that should be
+// periodically renewed, such as a database connection.
+//
+// Its Get method returns an instance of the resource, along with a release
+// function that the caller must invoke when it is done with the resource.
+//
+// The first call to Get creates a new resource instance. This instance is
+// cached and returned by subsequent calls to Get for a fixed duration. After
+// this duration expires, the next call to Get will create a new instance. The
+// expired instance is closed once all its users have released it.
+//
+// A Resource is safe for concurrent use.
+type Resource[T io.Closer] struct {
+	get      func() T
+	validFor time.Duration
+	after    func(func()) // for testing; fakes time.AfterFunc(validFor, f)
+
+	mu  sync.Mutex
+	cur *instance[T]
+}
+
+// New creates a new Resource that is valid for the given duration. The get
+// function is called to create a new resource instance when one is needed.
+func New[T io.Closer](get func() T, validFor time.Duration) *Resource[T] {
+	r := newAfter(get, func(f func()) {
+		time.AfterFunc(validFor, f)
+	})
+	r.validFor = validFor
+	return r
+}
+
+// newAfter returns a new resource with a fake after func, for testing.
+func newAfter[T io.Closer](get func() T, after func(func())) *Resource[T] {
+	return &Resource[T]{get: get, after: after}
+}
+
+// Get returns the current resource instance and a function to release it.
+// The release function must be called when the caller is done with the
+// resource.
+func (r *Resource[T]) Get() (T, func()) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if r.cur == nil {
+		r.cur = &instance[T]{v: r.get()}
+		// Acquire one ref for the new instance that lasts the duration.
+		//
+		// This is distinct from the ref acquired below; it ensures that the
+		// resource is not released until the duration has expired.
+		_, release := r.cur.acquire(true)
+		r.after(func() {
+			r.mu.Lock()
+			defer r.mu.Unlock()
+			release()
+			r.cur = nil
+		})
+	}
+	return r.cur.acquire(false)
+}