From 8baf09049ef743383ab0bf7e04bda93f321c56f7 Mon Sep 17 00:00:00 2001
From: ptrus <peter@u-s.si>
Date: Tue, 12 Apr 2022 16:15:48 +0200
Subject: [PATCH] worker/registration: fix waiting for deregister event

---
 .changelog/4662.bugfix.md                     |  1 +
 go/oasis-node/cmd/common/genesis/staking.go   |  3 +++
 .../e2e/runtime/keymanager_upgrade.go         | 20 +++++++++++++++++++
 go/worker/registration/worker.go              |  3 ---
 4 files changed, 24 insertions(+), 3 deletions(-)
 create mode 100644 .changelog/4662.bugfix.md

diff --git a/.changelog/4662.bugfix.md b/.changelog/4662.bugfix.md
new file mode 100644
index 00000000000..6a26bb84efb
--- /dev/null
+++ b/.changelog/4662.bugfix.md
@@ -0,0 +1 @@
+Fix waiting for deregister event on `RequestShutdown`
diff --git a/go/oasis-node/cmd/common/genesis/staking.go b/go/oasis-node/cmd/common/genesis/staking.go
index a6f4ba0f33a..d6da0cd2092 100644
--- a/go/oasis-node/cmd/common/genesis/staking.go
+++ b/go/oasis-node/cmd/common/genesis/staking.go
@@ -110,6 +110,9 @@ func NewAppendableStakingState() (*AppendableStakingState, error) {
 	st := &AppendableStakingState{
 		State: staking.Genesis{
 			Ledger: make(map[staking.Address]*staking.Account),
+			Parameters: staking.ConsensusParameters{
+				DebondingInterval: 1, // Minimum valid debonding interval.
+			},
 		},
 	}
 	if err := st.setDefaultFeeSplit(); err != nil {
diff --git a/go/oasis-test-runner/scenario/e2e/runtime/keymanager_upgrade.go b/go/oasis-test-runner/scenario/e2e/runtime/keymanager_upgrade.go
index 05d3f6a5d10..97fdf3ff244 100644
--- a/go/oasis-test-runner/scenario/e2e/runtime/keymanager_upgrade.go
+++ b/go/oasis-test-runner/scenario/e2e/runtime/keymanager_upgrade.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"fmt"
 	"path/filepath"
+	"time"
 
 	"github.com/oasisprotocol/oasis-core/go/common"
 	"github.com/oasisprotocol/oasis-core/go/common/cbor"
@@ -283,6 +284,12 @@ func (sc *kmUpgradeImpl) Run(childEnv *env.Env) error {
 		return err
 	}
 
+	nodeCh, nodeSub, err := sc.Net.Controller().Registry.WatchNodes(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to watch nodes: %w", err)
+	}
+	defer nodeSub.Close()
+
 	// Shutdown old keymanager and make sure it deregisters.
 	sc.Logger.Info("shutting down old keymanager")
 	oldKm := sc.Net.Keymanagers()[0]
@@ -290,6 +297,19 @@ func (sc *kmUpgradeImpl) Run(childEnv *env.Env) error {
 		return fmt.Errorf("failed to request shutdown: %w", err)
 	}
 
+	// Ensure keymanager deregisters.
+OUTER:
+	for {
+		select {
+		case ev := <-nodeCh:
+			if !ev.IsRegistration && ev.Node.ID.Equal(oldKm.NodeID) {
+				break OUTER
+			}
+		case <-time.After(10 * time.Second):
+			return fmt.Errorf("failed to wait for keymanager to de-register")
+		}
+	}
+
 	// Run client again.
 	sc.Logger.Info("starting a second client to check if key manager works")
 	newTestClient := sc.testClient.Clone().(*KeyValueEncTestClient)
diff --git a/go/worker/registration/worker.go b/go/worker/registration/worker.go
index 26a62af93c4..e151ee032f8 100644
--- a/go/worker/registration/worker.go
+++ b/go/worker/registration/worker.go
@@ -652,9 +652,6 @@ func (w *Worker) doNodeRegistration() {
 
 		case <-w.stopCh:
 			return
-
-		case <-w.stopRegCh:
-			return
 		}
 	}
 }