Merge pull request juju#18227 from Aflynn50/3.5-.36

juju#18227

Merge up 3.5 to 3.6. There were no merge conflicts. This contains:
- juju#18078 from juju
- juju#18170 from Aflynn50
- juju#18212 from gfouillet

<!-- 
The PR title should match: <type>(optional <scope>): <description>.

Please also ensure all commits in this PR comply with our conventional commits specification:
https://docs.google.com/document/d/1SYUo9G7qZ_jdoVXpUVamS5VCgHmtZ0QA-wZxKoMS-C0 
-->

state/modelmigration.go

@@ -849,10 +849,17 @@ func (st *State) LatestMigration() (ModelMigration, error) {
 	// away from a model and then migrated back.
 	if phase == migration.DONE {
 		model, err := st.Model()
-		if err != nil {
+		if errors.Is(err, errors.NotFound) {
+			// The model not being found breaks the precondition of this
+			// function, the model should be in state. However, to make this
+			// function more resilient and allow it to be called when the models
+			// existence is unknown, we do not return an error here.
+			logger.Debugf("checking latest migration: migrated model has been removed from the state")
+			return mig, nil
+		} else if err != nil {
 			return nil, errors.Trace(err)
 		}
-		if model.MigrationMode() == MigrationModeNone {
+		if model != nil && model.MigrationMode() == MigrationModeNone {
 			return nil, errors.NotFoundf("migration")
 		}
 	}

state/modelmigration_test.go

@@ -298,6 +298,8 @@ func (s *MigrationSuite) TestLatestMigrationWithPrevious(c *gc.C) {
 		migration.LOGTRANSFER,
 		migration.REAP,
 		migration.DONE,
+		// Check that it is idempotent on DONE.
+		migration.DONE,
 	}
 	for i := 0; i < 10; i++ {
 		mig, err := s.State2.CreateMigration(s.stdSpec)
@@ -311,6 +313,7 @@ func (s *MigrationSuite) TestLatestMigrationWithPrevious(c *gc.C) {
 	// Previous migration shouldn't be reported.
 	_, err := s.State2.LatestMigration()
 	c.Check(errors.IsNotFound(err), jc.IsTrue)
+	c.Check(err, gc.ErrorMatches, "migration not found")
 
 	// Start a new migration attempt, which should be reported.
 	migNext, err := s.State2.CreateMigration(s.stdSpec)
@@ -346,6 +349,11 @@ func (s *MigrationSuite) TestLatestRemovedModelMigration(c *gc.C) {
 	mig2, err := s.State2.CompletedMigration()
 	c.Assert(err, jc.ErrorIsNil)
 	c.Assert(mig2, jc.DeepEquals, mig1)
+
+	// Check that LatestMigration works with the model removed
+	mig3, err := s.State2.LatestMigration()
+	c.Assert(err, jc.ErrorIsNil)
+	c.Assert(mig3, jc.DeepEquals, mig1)
 }
 
 func (s *MigrationSuite) TestMigration(c *gc.C) {

worker/dbaccessor/tracker.go

@@ -138,15 +138,13 @@ func (w *trackedDBWorker) TxnNoRetry(ctx context.Context, fn func(context.Contex
 	// If the DB health check failed, the worker's error will be set,
 	// and we will be without a usable database reference. Return the error.
 	w.mutex.RLock()
+	defer w.mutex.RUnlock()
+
 	if w.err != nil {
-		w.mutex.RUnlock()
 		return errors.Trace(w.err)
 	}
 
-	db := w.db
-	w.mutex.RUnlock()
-
-	return errors.Trace(database.Txn(ctx, db, fn))
+	return errors.Trace(database.Txn(ctx, w.db, fn))
 }
 
 // meterDBOpResults decrements the active DB operation count,
@@ -200,7 +198,16 @@ func (w *trackedDBWorker) loop() error {
 			currentDB := w.db
 			w.mutex.RUnlock()
 
-			newDB, err := w.ensureDBAliveAndOpenIfRequired(currentDB)
+			ctx := w.tomb.Context(context.Background())
+			newDB, err := w.ensureDBAliveAndOpenIfRequired(ctx, currentDB)
+			if errors.Is(err, context.Canceled) {
+				select {
+				case <-w.tomb.Dying():
+					return tomb.ErrDying
+				default:
+					return errors.Trace(err)
+				}
+			}
 			if err != nil {
 				// If we get an error, ensure we close the underlying db and
 				// mark the tracked db in an error state.
@@ -240,11 +247,10 @@ func (w *trackedDBWorker) loop() error {
 // ensureDBAliveAndOpenNewIfRequired is a bit long-winded, but it is a way to
 // ensure that the underlying database is alive and well. If it is not, we
 // attempt to open a new one. If that fails, we return an error.
-func (w *trackedDBWorker) ensureDBAliveAndOpenIfRequired(db *sql.DB) (*sql.DB, error) {
+func (w *trackedDBWorker) ensureDBAliveAndOpenIfRequired(ctx context.Context, db *sql.DB) (*sql.DB, error) {
 	// Allow killing the tomb to cancel the context,
 	// so shutdown/restart can not be blocked by this call.
-	ctx, cancel := context.WithTimeout(context.TODO(), time.Second*10)
-	ctx = w.tomb.Context(ctx)
+	ctx, cancel := context.WithTimeout(ctx, time.Second*10)
 	defer cancel()
 
 	if w.logger.IsTraceEnabled() {

worker/dbaccessor/tracker_test.go

@@ -172,18 +172,15 @@ func (s *trackedDBWorkerSuite) TestWorkerAttemptsToVerifyDBButSucceeds(c *gc.C)
 
 	s.timer.EXPECT().Reset(PollInterval).Times(1)
 
-	dbChange := make(chan struct{})
-	s.dbApp.EXPECT().Open(gomock.Any(), "controller").Return(s.DB(), nil).Times(DefaultVerifyAttempts - 1)
-	s.dbApp.EXPECT().Open(gomock.Any(), "controller").Return(s.DB(), nil).DoAndReturn(func(_ context.Context, _ string) (*sql.DB, error) {
-		defer close(dbChange)
-		return s.DB(), nil
-	})
+	dbReady := make(chan struct{})
+	s.dbApp.EXPECT().Open(gomock.Any(), "controller").Return(s.DB(), nil).Times(DefaultVerifyAttempts)
 
 	var count uint64
 	pingFn := func(context.Context, *sql.DB) error {
 		val := atomic.AddUint64(&count, 1)
 
 		if val == DefaultVerifyAttempts {
+			defer close(dbReady)
 			return nil
 		}
 		return errors.New("boom")
@@ -200,9 +197,9 @@ func (s *trackedDBWorkerSuite) TestWorkerAttemptsToVerifyDBButSucceeds(c *gc.C)
 		c.Fatal("timed out waiting for DB callback")
 	}
 
-	// The db should have changed to the new db.
+	// The db should wait to a successful ping after several attempts
 	select {
-	case <-dbChange:
+	case <-dbReady:
 	case <-time.After(testing.ShortWait):
 		c.Fatal("timed out waiting for DB callback")
 	}