From 4d4fc4dbd9083ff1ff32fa864359131ccd29922a Mon Sep 17 00:00:00 2001
From: Chris Goller <goller@gmail.com>
Date: Wed, 30 Aug 2023 17:47:24 -0500
Subject: [PATCH] fix(boltdb): close cache and history dbs on exit

Signed-off-by: Chris Goller <goller@gmail.com>
---
 cmd/buildkitd/main.go               |  1 +
 control/control.go                  | 14 +++++++++++++-
 solver/bboltcachestorage/storage.go |  4 ++++
 3 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/cmd/buildkitd/main.go b/cmd/buildkitd/main.go
index 7a245a2c6..25a4bb431 100644
--- a/cmd/buildkitd/main.go
+++ b/cmd/buildkitd/main.go
@@ -754,6 +754,7 @@ func newController(c *cli.Context, cfg *config.Config) (*control.Controller, err
 		Entitlements:              cfg.Entitlements,
 		TraceCollector:            tc,
 		HistoryDB:                 historyDB,
+		CacheStore:                cacheStorage,
 		LeaseManager:              w.LeaseManager(),
 		ContentStore:              w.ContentStore(),
 		HistoryConfig:             cfg.History,
diff --git a/control/control.go b/control/control.go
index 8a09b8ace..1afd518a6 100644
--- a/control/control.go
+++ b/control/control.go
@@ -12,6 +12,7 @@ import (
 	"github.com/containerd/containerd/content"
 	"github.com/containerd/containerd/services/content/contentserver"
 	"github.com/docker/distribution/reference"
+	"github.com/hashicorp/go-multierror"
 	"github.com/mitchellh/hashstructure/v2"
 	controlapi "github.com/moby/buildkit/api/services/control"
 	apitypes "github.com/moby/buildkit/api/types"
@@ -29,6 +30,7 @@ import (
 	"github.com/moby/buildkit/session/grpchijack"
 	containerdsnapshot "github.com/moby/buildkit/snapshot/containerd"
 	"github.com/moby/buildkit/solver"
+	"github.com/moby/buildkit/solver/bboltcachestorage"
 	"github.com/moby/buildkit/solver/llbsolver"
 	"github.com/moby/buildkit/solver/llbsolver/proc"
 	"github.com/moby/buildkit/solver/pb"
@@ -61,6 +63,7 @@ type Opt struct {
 	Entitlements              []string
 	TraceCollector            sdktrace.SpanExporter
 	HistoryDB                 *bbolt.DB
+	CacheStore                *bboltcachestorage.Store
 	LeaseManager              *leaseutil.Manager
 	ContentStore              *containerdsnapshot.Store
 	HistoryConfig             *config.HistoryConfig
@@ -123,7 +126,16 @@ func NewController(opt Opt) (*Controller, error) {
 }
 
 func (c *Controller) Close() error {
-	return c.opt.WorkerController.Close()
+	rerr := c.opt.HistoryDB.Close()
+	if err := c.opt.WorkerController.Close(); err != nil {
+		rerr = multierror.Append(rerr, err)
+	}
+
+	if err := c.opt.CacheStore.Close(); err != nil {
+		rerr = multierror.Append(rerr, err)
+	}
+
+	return rerr
 }
 
 func (c *Controller) Register(server *grpc.Server) {
diff --git a/solver/bboltcachestorage/storage.go b/solver/bboltcachestorage/storage.go
index 515feffbf..37feb03a8 100644
--- a/solver/bboltcachestorage/storage.go
+++ b/solver/bboltcachestorage/storage.go
@@ -54,6 +54,10 @@ func (s *Store) Exists(id string) bool {
 	return exists
 }
 
+func (s *Store) Close() error {
+	return s.db.Close()
+}
+
 func (s *Store) Walk(fn func(id string) error) error {
 	ids := make([]string, 0)
 	if err := s.db.View(func(tx *bolt.Tx) error {