From 6d7f78b33d0b801f65b9ef0233ef913335522d1e Mon Sep 17 00:00:00 2001
From: iofq <cjriddz@protonmail.com>
Date: Mon, 2 Sep 2024 18:58:36 -0500
Subject: [PATCH] fix(api): Clean up database locks on SIGTERM

---
 api/mirror.go        | 15 ++++++++++-----
 cmd/mirror_update.go | 12 ++++++++----
 context/context.go   |  3 ++-
 3 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/api/mirror.go b/api/mirror.go
index 4c4adf7fc..b78716dfb 100644
--- a/api/mirror.go
+++ b/api/mirror.go
@@ -393,13 +393,20 @@ func apiMirrorsUpdate(c *gin.Context) {
 			return &task.ProcessReturnValue{Code: http.StatusInternalServerError, Value: nil}, fmt.Errorf("unable to update: %s", err)
 		}
 
-		defer func() {
+		cleanupDB := func() {
 			// on any interruption, unlock the mirror
-			e := context.ReOpenDatabase()
-			if e == nil {
+			err := context.ReOpenDatabase()
+			if err == nil {
 				remote.MarkAsIdle()
 				collection.Update(remote)
 			}
+		}
+		defer cleanupDB()
+
+		context.GoContextHandleSignals()
+		go func() {
+			<-context.Done()
+			cleanupDB()
 		}()
 
 		remote.MarkAsUpdating()
@@ -408,8 +415,6 @@ func apiMirrorsUpdate(c *gin.Context) {
 			return &task.ProcessReturnValue{Code: http.StatusInternalServerError, Value: nil}, fmt.Errorf("unable to update: %s", err)
 		}
 
-		context.GoContextHandleSignals()
-
 		count := len(queue)
 		taskDetail := struct {
 			TotalDownloadSize         int64
diff --git a/cmd/mirror_update.go b/cmd/mirror_update.go
index 2e6df4397..1e41960b6 100644
--- a/cmd/mirror_update.go
+++ b/cmd/mirror_update.go
@@ -91,18 +91,24 @@ func aptlyMirrorUpdate(cmd *commander.Command, args []string) error {
 	context.Progress().Printf("Building download queue...\n")
 	queue, downloadSize, err = repo.BuildDownloadQueue(context.PackagePool(), collectionFactory.PackageCollection(),
 		collectionFactory.ChecksumCollection(nil), skipExistingPackages)
-
 	if err != nil {
 		return fmt.Errorf("unable to update: %s", err)
 	}
 
-	defer func() {
+	cleanupDB := func() {
 		// on any interruption, unlock the mirror
 		err = context.ReOpenDatabase()
 		if err == nil {
 			repo.MarkAsIdle()
 			collectionFactory.RemoteRepoCollection().Update(repo)
 		}
+	}
+	defer cleanupDB()
+
+	context.GoContextHandleSignals()
+	go func() {
+		<-context.Done()
+		cleanupDB()
 	}()
 
 	repo.MarkAsUpdating()
@@ -116,8 +122,6 @@ func aptlyMirrorUpdate(cmd *commander.Command, args []string) error {
 		return fmt.Errorf("unable to update: %s", err)
 	}
 
-	context.GoContextHandleSignals()
-
 	count := len(queue)
 	context.Progress().Printf("Download queue: %d items (%s)\n", count, utils.HumanBytes(downloadSize))
 
diff --git a/context/context.go b/context/context.go
index 7cc7c6efd..fa48464c1 100644
--- a/context/context.go
+++ b/context/context.go
@@ -13,6 +13,7 @@ import (
 	"runtime/pprof"
 	"strings"
 	"sync"
+	"syscall"
 	"time"
 
 	"github.com/aptly-dev/aptly/aptly"
@@ -558,7 +559,7 @@ func (context *AptlyContext) GoContextHandleSignals() {
 
 	// Catch ^C
 	sigch := make(chan os.Signal, 1)
-	signal.Notify(sigch, os.Interrupt)
+	signal.Notify(sigch, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
 
 	var cancel gocontext.CancelFunc