From 0b7b585ba061ad85c4f16c3180522791740d6538 Mon Sep 17 00:00:00 2001
From: Ales Stimec <ales.stimec@canonical.com>
Date: Mon, 24 Jun 2024 10:11:17 +0200
Subject: [PATCH] Fixes updates to controller's UnavailableSince field.

---
 internal/jimm/watcher.go      | 54 +++++++++++++-------------
 internal/jimm/watcher_test.go | 73 ++++++++++++++++++++++++++++++++++-
 2 files changed, 97 insertions(+), 30 deletions(-)

diff --git a/internal/jimm/watcher.go b/internal/jimm/watcher.go
index 802917b55..cb4a8e5a3 100644
--- a/internal/jimm/watcher.go
+++ b/internal/jimm/watcher.go
@@ -130,27 +130,38 @@ func (w *Watcher) WatchAllModelSummaries(ctx context.Context, interval time.Dura
 	}
 }
 
-func (w *Watcher) dialController(ctx context.Context, ctl *dbmodel.Controller) (API, error) {
+func (w *Watcher) dialController(ctx context.Context, ctl *dbmodel.Controller) (api API, err error) {
 	const op = errors.Op("jimm.dialController")
 
-	// connect to the controller
-	api, err := w.Dialer.Dial(ctx, ctl, names.ModelTag{}, nil)
-	if err != nil {
-		if !ctl.UnavailableSince.Valid {
-			ctl.UnavailableSince = db.Now()
-			var err error
-			if err = w.Database.UpdateController(ctx, ctl); err != nil {
-				zapctx.Error(ctx, "cannot set controller unavailable", zap.Error(err))
-			}
-			if w.controllerUnavailableChan != nil {
-				select {
-				case w.controllerUnavailableChan <- err:
-				default:
-				}
+	updateController := false
+	defer func() {
+		if !updateController {
+			return
+		}
+		if uerr := w.Database.UpdateController(ctx, ctl); err != nil {
+			zapctx.Error(ctx, "cannot set controller available", zap.Error(uerr))
+		}
+		// Note (alesstimec) This channel is only available in tests.
+		if w.controllerUnavailableChan != nil {
+			select {
+			case w.controllerUnavailableChan <- err:
+			default:
 			}
 		}
+	}()
+
+	// connect to the controller
+	api, err = w.Dialer.Dial(ctx, ctl, names.ModelTag{}, nil)
+	if err != nil {
+		ctl.UnavailableSince = db.Now()
+		updateController = true
+
 		return nil, errors.E(op, err)
 	}
+	if ctl.UnavailableSince.Valid {
+		ctl.UnavailableSince = sql.NullTime{}
+		updateController = true
+	}
 	return api, nil
 }
 
@@ -350,19 +361,6 @@ func (w *Watcher) watchAllModelSummaries(ctx context.Context, ctl *dbmodel.Contr
 	// connect to the controller
 	api, err := w.dialController(ctx, ctl)
 	if err != nil {
-		if !ctl.UnavailableSince.Valid {
-			ctl.UnavailableSince = db.Now()
-			var err error
-			if err = w.Database.UpdateController(ctx, ctl); err != nil {
-				zapctx.Error(ctx, "cannot set controller unavailable", zap.Error(err))
-			}
-			if w.controllerUnavailableChan != nil {
-				select {
-				case w.controllerUnavailableChan <- err:
-				default:
-				}
-			}
-		}
 		return errors.E(op, err)
 	}
 	defer api.Close()
diff --git a/internal/jimm/watcher_test.go b/internal/jimm/watcher_test.go
index 2fcbfd962..7f81a8575 100644
--- a/internal/jimm/watcher_test.go
+++ b/internal/jimm/watcher_test.go
@@ -805,13 +805,82 @@ func TestWatcherSetsControllerUnavailable(t *testing.T) {
 		err = w.Database.GetController(ctx, &ctl)
 		c.Assert(err, qt.IsNil)
 		c.Check(ctl.UnavailableSince.Valid, qt.Equals, true)
-		c.Logf("%v %v", ctl.UnavailableSince.Time, time.Now())
-		c.Check(ctl.UnavailableSince.Time.After(time.Now().Add(-10*time.Millisecond)), qt.Equals, true)
 	}
 	cancel()
 	wg.Wait()
 }
 
+func TestWatcherClearsControllerUnavailable(t *testing.T) {
+	c := qt.New(t)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	w := jimm.Watcher{
+		Database: db.Database{
+			DB: jimmtest.PostgresDB(c, nil),
+		},
+		Dialer: &jimmtest.Dialer{
+			API: &jimmtest.API{
+				AllModelWatcherNext_: func(_ context.Context, _ string) ([]jujuparams.Delta, error) {
+					cancel()
+					<-ctx.Done()
+					return nil, ctx.Err()
+				},
+				ModelInfo_: func(_ context.Context, info *jujuparams.ModelInfo) error {
+					switch info.UUID {
+					default:
+						c.Errorf("unexpected model uuid: %s", info.UUID)
+					case "00000002-0000-0000-0000-000000000002":
+					case "00000002-0000-0000-0000-000000000003":
+					}
+					return errors.E(errors.CodeNotFound)
+				},
+				WatchAllModels_: func(ctx context.Context) (string, error) {
+					return "1234", nil
+				},
+			},
+		},
+		Pubsub: &testPublisher{},
+	}
+
+	env := jimmtest.ParseEnvironment(c, testWatcherEnv)
+	err := w.Database.Migrate(ctx, false)
+	c.Assert(err, qt.IsNil)
+	env.PopulateDB(c, w.Database)
+
+	// update controller's UnavailableSince field
+	ctl := dbmodel.Controller{
+		Name: "controller-1",
+	}
+	err = w.Database.GetController(ctx, &ctl)
+	c.Assert(err, qt.IsNil)
+	ctl.UnavailableSince = sql.NullTime{
+		Time:  time.Now(),
+		Valid: true,
+	}
+	err = w.Database.UpdateController(ctx, &ctl)
+	c.Assert(err, qt.IsNil)
+
+	// start the watcher
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		err := w.Watch(ctx, time.Millisecond)
+		checkIfContextCanceled(c, ctx, err)
+	}()
+	wg.Wait()
+
+	// check that the unavailable since time has been cleared
+	ctl = dbmodel.Controller{
+		Name: "controller-1",
+	}
+	err = w.Database.GetController(context.Background(), &ctl)
+	c.Assert(err, qt.IsNil)
+	c.Assert(ctl.UnavailableSince.Valid, qt.IsFalse)
+}
+
 func TestWatcherRemoveDyingModelsOnStartup(t *testing.T) {
 	c := qt.New(t)