From a133ab31fa1289badc3c77f676b7049949dead4e Mon Sep 17 00:00:00 2001
From: Dmitry Verkhoturov <paskal.07@gmail.com>
Date: Sun, 12 Apr 2020 10:04:27 +0200
Subject: [PATCH] add context cancellation to image.Service.Cancel()

---
 backend/app/cmd/server.go             |  4 +++-
 backend/app/store/image/image.go      | 10 ++++++++--
 backend/app/store/image/image_test.go |  4 ++--
 3 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/backend/app/cmd/server.go b/backend/app/cmd/server.go
index b073a671..fb488d11 100644
--- a/backend/app/cmd/server.go
+++ b/backend/app/cmd/server.go
@@ -502,7 +502,9 @@ func (a *serverApp) run(ctx context.Context) error {
 		log.Printf("[WARN] failed to close avatar store, %s", e)
 	}
 	a.notifyService.Close()
-	a.imageService.Close()
+	// call potentially infinite loop with cancellation after a minute as a safeguard
+	minuteCtx, _ := context.WithTimeout(context.Background(), time.Minute)
+	a.imageService.Close(minuteCtx)
 
 	close(a.terminated)
 	return nil
diff --git a/backend/app/store/image/image.go b/backend/app/store/image/image.go
index 4ef37687..638b5928 100644
--- a/backend/app/store/image/image.go
+++ b/backend/app/store/image/image.go
@@ -149,14 +149,20 @@ func (s *Service) Cleanup(ctx context.Context) {
 }
 
 // Close flushes all in-progress submits and enforces waiting commits
-func (s *Service) Close() {
+func (s *Service) Close(ctx context.Context) {
 	log.Printf("[INFO] close image service ")
 	// terminate commit goroutine using s.term only if it is started
 	if atomic.LoadInt32(&s.term) == 1 {
 		atomic.StoreInt32(&s.term, 0) // enforce non-delayed commits for all ids left in submitCh
+		var ctxCancel bool
 		for {
+			select {
+			case <-ctx.Done():
+				ctxCancel = true
+			default:
+			}
 			// set to 1 by commit goroutine after everything waited on TTL sent
-			if atomic.LoadInt32(&s.term) == 1 {
+			if atomic.LoadInt32(&s.term) == 1 || ctxCancel {
 				break
 			}
 			time.Sleep(10 * time.Millisecond)
diff --git a/backend/app/store/image/image_test.go b/backend/app/store/image/image_test.go
index 27ed6451..7afb4251 100644
--- a/backend/app/store/image/image_test.go
+++ b/backend/app/store/image/image_test.go
@@ -128,7 +128,7 @@ func TestService_Close(t *testing.T) {
 	svc.Submit(func() []string { return []string{"id1", "id2", "id3"} })
 	svc.Submit(func() []string { return []string{"id4", "id5"} })
 	svc.Submit(nil)
-	svc.Close()
+	svc.Close(context.TODO())
 	store.AssertNumberOfCalls(t, "Commit", 5)
 }
 
@@ -141,7 +141,7 @@ func TestService_SubmitDelay(t *testing.T) {
 	svc.Submit(func() []string { return []string{"id4", "id5"} })
 	svc.Submit(nil)
 	store.AssertNumberOfCalls(t, "Commit", 3)
-	svc.Close()
+	svc.Close(context.TODO())
 	store.AssertNumberOfCalls(t, "Commit", 5)
 }