add context cancellation to backup routines

4 files changed, 75 insertions, 29 deletions

diff --git a/system/admin/handlers.go b/system/admin/handlers.go
index 4f8ae83..4734ba0 100644
--- a/system/admin/handlers.go
+++ b/system/admin/handlers.go
@@ -195,9 +195,12 @@ func configHandler(res http.ResponseWriter, req *http.Request) {
 }
 
 func backupHandler(res http.ResponseWriter, req *http.Request) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	switch req.URL.Query().Get("source") {
 	case "system":
-		err := db.Backup(res)
+		err := db.Backup(ctx, res)
 		if err != nil {
 			log.Println("Failed to run backup on system:", err)
 			res.WriteHeader(http.StatusInternalServerError)
@@ -205,7 +208,7 @@ func backupHandler(res http.ResponseWriter, req *http.Request) {
 		}
 
 	case "analytics":
-		err := analytics.Backup(res)
+		err := analytics.Backup(ctx, res)
 		if err != nil {
 			log.Println("Failed to run backup on analytics:", err)
 			res.WriteHeader(http.StatusInternalServerError)
@@ -213,7 +216,7 @@ func backupHandler(res http.ResponseWriter, req *http.Request) {
 		}
 
 	case "uploads":
-		err := upload.Backup(res)
+		err := upload.Backup(ctx, res)
 		if err != nil {
 			log.Println("Failed to run backup on uploads:", err)
 			res.WriteHeader(http.StatusInternalServerError)
diff --git a/system/admin/upload/backup.go b/system/admin/upload/backup.go
index 28b1b8e..b4f6393 100644
--- a/system/admin/upload/backup.go
+++ b/system/admin/upload/backup.go
@@ -3,6 +3,7 @@ package upload
 import (
 	"archive/tar"
 	"compress/gzip"
+	"context"
 	"fmt"
 	"io"
 	"net/http"
@@ -13,7 +14,7 @@ import (
 
 // Backup creates an archive of a project's uploads and writes it
 // to the response as a download
-func Backup(res http.ResponseWriter) error {
+func Backup(ctx context.Context, res http.ResponseWriter) error {
 	ts := time.Now().Unix()
 	filename := fmt.Sprintf("uploads-%d.bak.tar.gz", ts)
 	tmp := os.TempDir()
@@ -30,7 +31,8 @@ func Backup(res http.ResponseWriter) error {
 	gz := gzip.NewWriter(f)
 	tarball := tar.NewWriter(gz)
 
-	err = filepath.Walk("uploads", func(path string, info os.FileInfo, err error) error {
+	errChan := make(chan error, 1)
+	walkFn := func(path string, info os.FileInfo, err error) error {
 		if err != nil {
 			return err
 		}
@@ -53,6 +55,7 @@ func Backup(res http.ResponseWriter) error {
 				return err
 			}
 			defer src.Close()
+
 			_, err = io.Copy(tarball, src)
 			if err != nil {
 				return err
@@ -70,6 +73,24 @@ func Backup(res http.ResponseWriter) error {
 		}
 
 		return nil
+	}
+
+	// stop processing if we get a cancellation signal
+	err = filepath.Walk("uploads", func(path string, info os.FileInfo, err error) error {
+		go func() { errChan <- walkFn(path, info, err) }()
+
+		select {
+		case <-ctx.Done():
+			if err := ctx.Err(); err != nil {
+				return err
+			}
+		case err := <-errChan:
+			if err != nil {
+				return err
+			}
+		}
+
+		return nil
 	})
 	if err != nil {
 		fmt.Println(err)
diff --git a/system/api/analytics/backup.go b/system/api/analytics/backup.go
index 07b1a46..46261a1 100644
--- a/system/api/analytics/backup.go
+++ b/system/api/analytics/backup.go
@@ -1,6 +1,7 @@
 package analytics
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"time"
@@ -8,19 +9,29 @@ import (
 	"github.com/boltdb/bolt"
 )
 
-// Backup writes a snapshot of the analytics.db database to an HTTP response
-func Backup(res http.ResponseWriter) error {
-	err := store.View(func(tx *bolt.Tx) error {
-		ts := time.Now().Unix()
-		disposition := `attachment; filename="analytics-%d.db.bak"`
+// Backup writes a snapshot of the system.db database to an HTTP response. The
+// output is discarded if we get a cancellation signal.
+func Backup(ctx context.Context, res http.ResponseWriter) error {
+	errChan := make(chan error, 1)
 
-		res.Header().Set("Content-Type", "application/octet-stream")
-		res.Header().Set("Content-Disposition", fmt.Sprintf(disposition, ts))
-		res.Header().Set("Content-Length", fmt.Sprintf("%d", int(tx.Size())))
+	go func() {
+		errChan <- store.View(func(tx *bolt.Tx) error {
+			ts := time.Now().Unix()
+			disposition := `attachment; filename="analytics-%d.db.bak"`
 
-		_, err := tx.WriteTo(res)
-		return err
-	})
+			res.Header().Set("Content-Type", "application/octet-stream")
+			res.Header().Set("Content-Disposition", fmt.Sprintf(disposition, ts))
+			res.Header().Set("Content-Length", fmt.Sprintf("%d", int(tx.Size())))
+
+			_, err := tx.WriteTo(res)
+			return err
+		})
+	}()
 
-	return err
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case err := <-errChan:
+		return err
+	}
 }
diff --git a/system/db/backup.go b/system/db/backup.go
index 735abe4..4a2bf96 100644
--- a/system/db/backup.go
+++ b/system/db/backup.go
@@ -1,6 +1,7 @@
 package db
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"time"
@@ -8,19 +9,29 @@ import (
 	"github.com/boltdb/bolt"
 )
 
-// Backup writes a snapshot of the system.db database to an HTTP response
-func Backup(res http.ResponseWriter) error {
-	err := store.View(func(tx *bolt.Tx) error {
-		ts := time.Now().Unix()
-		disposition := `attachment; filename="system-%d.db.bak"`
+// Backup writes a snapshot of the system.db database to an HTTP response. The
+// output is discarded if we get a cancellation signal.
+func Backup(ctx context.Context, res http.ResponseWriter) error {
+	errChan := make(chan error, 1)
 
-		res.Header().Set("Content-Type", "application/octet-stream")
-		res.Header().Set("Content-Disposition", fmt.Sprintf(disposition, ts))
-		res.Header().Set("Content-Length", fmt.Sprintf("%d", int(tx.Size())))
+	go func() {
+		errChan <- store.View(func(tx *bolt.Tx) error {
+			ts := time.Now().Unix()
+			disposition := `attachment; filename="system-%d.db.bak"`
 
-		_, err := tx.WriteTo(res)
-		return err
-	})
+			res.Header().Set("Content-Type", "application/octet-stream")
+			res.Header().Set("Content-Disposition", fmt.Sprintf(disposition, ts))
+			res.Header().Set("Content-Length", fmt.Sprintf("%d", int(tx.Size())))
+
+			_, err := tx.WriteTo(res)
+			return err
+		})
+	}()
 
-	return err
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case err := <-errChan:
+		return err
+	}
 }