VAULT-8631 Make upgrade synchronous when no keys to upgrade (#66)

backend.go

@@ -90,7 +90,7 @@ func Factory(ctx context.Context, conf *logical.BackendConfig) (logical.Backend,
 	return b, nil
 }
 
-// Factory returns a new backend as logical.Backend.
+// VersionedKVFactory returns a new KVV2 backend as logical.Backend.
 func VersionedKVFactory(ctx context.Context, conf *logical.BackendConfig) (logical.Backend, error) {
 	upgradeCtx, upgradeCancelFunc := context.WithCancel(ctx)
 

path_config_test.go

@@ -157,5 +157,4 @@ func TestVersionedKV_Config_DeleteVersionAfter(t *testing.T) {
 			}
 		})
 	}
-
 }

path_data_test.go

@@ -30,36 +30,23 @@ func getBackend(t *testing.T) (logical.Backend, logical.Storage) {
 		t.Fatalf("unable to create backend: %v", err)
 	}
 
-	// Wait for the upgrade to finish
-	timeout := time.After(20 * time.Second)
-	ticker := time.Tick(time.Second)
-
-	for {
-		select {
-		case <-timeout:
-			t.Fatal("timeout expired waiting for upgrade")
-		case <-ticker:
-			req := &logical.Request{
-				Operation: logical.ReadOperation,
-				Path:      "config",
-				Storage:   config.StorageView,
-			}
-
-			resp, err := b.HandleRequest(context.Background(), req)
-			if err != nil {
-				t.Fatalf("unable to read config: %s", err.Error())
-				return nil, nil
-			}
-
-			if resp != nil && !resp.IsError() {
-				return b, config.StorageView
-			}
-
-			if resp == nil || (resp.IsError() && strings.Contains(resp.Error().Error(), "Upgrading from non-versioned to versioned")) {
-				t.Log("waiting for upgrade to complete")
-			}
-		}
+	req := &logical.Request{
+		Operation: logical.ReadOperation,
+		Path:      "config",
+		Storage:   config.StorageView,
 	}
+
+	resp, err := b.HandleRequest(context.Background(), req)
+	if err != nil {
+		t.Fatalf("unable to read config: %s", err.Error())
+		return nil, nil
+	}
+
+	if resp == nil || resp.IsError() {
+		t.Fatalf("Error during mount creation: %x", resp.Error().Error())
+	}
+
+	return b, config.StorageView
 }
 
 // getKeySet will produce a set of the keys that exist in m

upgrade.go

@@ -118,6 +118,18 @@ func (b *versionedKVBackend) Upgrade(ctx context.Context, s logical.Storage) err
 		return nil
 	}
 
+	// If we have 0 keys, it's either a new mount or one that's trivial to upgrade,
+	// so we should do the upgrade synchronously
+	upgradeSynchronously := false
+	keys, err := logical.CollectKeys(ctx, s)
+	if err != nil {
+		b.Logger().Error("upgrading resulted in error", "error", err)
+		return err
+	}
+	if len(keys) == 0 {
+		upgradeSynchronously = true
+	}
+
 	upgradeInfo := &UpgradeInfo{
 		StartedTime: ptypes.TimestampNow(),
 	}
@@ -128,7 +140,7 @@ func (b *versionedKVBackend) Upgrade(ctx context.Context, s logical.Storage) err
 		return err
 	}
 
-	// Because this is a long running process we need a new context.
+	// Because this is a long-running process we need a new context.
 	ctx = context.Background()
 
 	upgradeKey := func(key string) error {
@@ -189,10 +201,7 @@ func (b *versionedKVBackend) Upgrade(ctx context.Context, s logical.Storage) err
 		return nil
 	}
 
-	// Run the actual upgrade in a go routine so we don't block the client on a
-	// potentially long process.
-	go func() {
-
+	upgradeFunc := func() {
 		// Write the canary value and if we are read only wait until the setup
 		// process has finished.
 	READONLY_LOOP:
@@ -257,7 +266,15 @@ func (b *versionedKVBackend) Upgrade(ctx context.Context, s logical.Storage) err
 		}
 
 		atomic.StoreUint32(b.upgrading, 0)
-	}()
+	}
+
+	if upgradeSynchronously {
+		upgradeFunc()
+	} else {
+		// We run the actual upgrade in a go routine, so we don't block the client on a
+		// potentially long process.
+		go upgradeFunc()
+	}
 
 	return nil
 }