From 48f2c56caa5b42fdd33b5f4298478e48c5b9f961 Mon Sep 17 00:00:00 2001
From: James Houlahan <james.houlahan@proton.ch>
Date: Tue, 7 Feb 2023 12:29:04 +0100
Subject: [PATCH] fix(GODT-2327): Better sleep (with context)

---
 internal/bridge/bridge.go |  2 --
 internal/bridge/user.go   |  5 +++--
 internal/user/user.go     | 14 +++++++++++---
 3 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/internal/bridge/bridge.go b/internal/bridge/bridge.go
index 886c740b..d3ea7d07 100644
--- a/internal/bridge/bridge.go
+++ b/internal/bridge/bridge.go
@@ -376,8 +376,6 @@ func (bridge *Bridge) init(tlsReporter TLSReporter) error {
 
 	// Attempt to lazy load users when triggered.
 	bridge.goLoad = bridge.tasks.Trigger(func(ctx context.Context) {
-		logrus.Info("Loading users")
-
 		if err := bridge.loadUsers(ctx); err != nil {
 			logrus.WithError(err).Error("Failed to load users")
 		} else {
diff --git a/internal/bridge/user.go b/internal/bridge/user.go
index 25a697b7..1a250365 100644
--- a/internal/bridge/user.go
+++ b/internal/bridge/user.go
@@ -330,17 +330,18 @@ func (bridge *Bridge) loginUser(ctx context.Context, client *proton.Client, auth
 // loadUsers tries to load each user in the vault that isn't already loaded.
 func (bridge *Bridge) loadUsers(ctx context.Context) error {
 	logrus.WithField("count", len(bridge.vault.GetUserIDs())).Info("Loading users")
+	defer logrus.Info("Finished loading users")
 
 	return bridge.vault.ForUser(runtime.NumCPU(), func(user *vault.User) error {
 		log := logrus.WithField("userID", user.UserID())
 
 		if user.AuthUID() == "" {
-			log.Info("Not loading disconnected user")
+			log.Info("User is not connected (skipping)")
 			return nil
 		}
 
 		if safe.RLockRet(func() bool { return mapHas(bridge.users, user.UserID()) }, bridge.usersLock) {
-			log.Debug("Not loading already-loaded user")
+			log.Info("User is already loaded (skipping)")
 			return nil
 		}
 
diff --git a/internal/user/user.go b/internal/user/user.go
index 58bfcd09..53a8136c 100644
--- a/internal/user/user.go
+++ b/internal/user/user.go
@@ -185,12 +185,12 @@ func New(
 
 	// When triggered, sync the user and then begin streaming API events.
 	user.goSync = user.tasks.Trigger(func(ctx context.Context) {
-		user.log.Debug("Sync triggered")
+		user.log.Info("Sync triggered")
 
 		// Sync the user.
 		user.syncAbort.Do(ctx, func(ctx context.Context) {
 			if user.vault.SyncStatus().IsComplete() {
-				user.log.Debug("Sync already complete, skipping")
+				user.log.Info("Sync already complete, skipping")
 				return
 			}
 
@@ -200,7 +200,7 @@ func New(
 					return
 				} else if err := user.doSync(ctx); err != nil {
 					user.log.WithError(err).Error("Failed to sync, will retry later")
-					time.Sleep(SyncRetryCooldown)
+					sleepCtx(ctx, SyncRetryCooldown)
 				} else {
 					user.log.Info("Sync complete, starting API event stream")
 					return
@@ -685,3 +685,11 @@ func b32(b bool) uint32 {
 
 	return 0
 }
+
+// sleepCtx sleeps for the given duration, or until the context is canceled.
+func sleepCtx(ctx context.Context, d time.Duration) {
+	select {
+	case <-ctx.Done():
+	case <-time.After(d):
+	}
+}