From dece3bd899db4e81e8b3edbd6b0ae4cf3702e90a Mon Sep 17 00:00:00 2001
From: gram-signal <84339875+gram-signal@users.noreply.github.com>
Date: Thu, 14 Nov 2024 12:16:03 -0800
Subject: [PATCH] Add enforced timeout to Raft join.

---
 host/config/config.go           |  1 +
 host/config/raft.go             |  2 ++
 host/raftmanager/raftmanager.go | 23 ++++++++++++++++-------
 host/service/service.go         |  2 +-
 4 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/host/config/config.go b/host/config/config.go
index 8eda50b..1286656 100644
--- a/host/config/config.go
+++ b/host/config/config.go
@@ -114,6 +114,7 @@ func Default() *Config {
 			EnvStatsPollDuration:       0,
 			RefreshAttestationDuration: time.Minute * 10,
 			EnclaveConcurrency:         util.Min(runtime.NumCPU(), 64),
+			InitialJoinDuration:        time.Hour,
 		},
 		Redis: RedisConfig{
 			Name:             "test",
diff --git a/host/config/raft.go b/host/config/raft.go
index 5587bf3..132a70e 100644
--- a/host/config/raft.go
+++ b/host/config/raft.go
@@ -21,6 +21,8 @@ type RaftHostConfig struct {
 	EnvStatsPollDuration time.Duration `yaml:"envStatsPollDuration"`
 	// max number of in-flight enclave calls
 	EnclaveConcurrency int `yaml:"enclaveConcurrency"`
+	// timeout for intitial raft Joining attempt
+	InitialJoinDuration time.Duration `yaml:"initialJoinDuration"`
 }
 
 func (r *RaftHostConfig) validate() []string {
diff --git a/host/raftmanager/raftmanager.go b/host/raftmanager/raftmanager.go
index ef194e9..d6829ee 100644
--- a/host/raftmanager/raftmanager.go
+++ b/host/raftmanager/raftmanager.go
@@ -132,14 +132,23 @@ func (r *RaftManager) CreateOrJoin(ctx context.Context) error {
 	if err != nil {
 		return errors.New("failed to fetch raft peers")
 	}
-	if raftPeer == r.me {
-		logger.Infow("attempting to create a new raft group")
-		if err := r.createRaft(); err != nil {
-			return err
+	joined := make(chan error)
+
+	go func() {
+		if raftPeer == r.me {
+			logger.Infow("attempting to create a new raft group")
+			joined <- r.createRaft()
+		} else {
+			logger.Infow("attempting to join existing raft group", "peerID", raftPeer)
+			joined <- r.joinExistingRaftPeer(raftPeer)
 		}
-	} else {
-		logger.Infow("attempting to join existing raft group", "peerID", raftPeer)
-		if err := r.joinExistingRaftPeer(raftPeer); err != nil {
+	}()
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case err := <-joined:
+		if err != nil {
 			return err
 		}
 	}
diff --git a/host/service/service.go b/host/service/service.go
index e78be32..e14dcd1 100644
--- a/host/service/service.go
+++ b/host/service/service.go
@@ -118,7 +118,7 @@ func Start(ctx context.Context, hconfig *config.Config, authenticator auth.Auth,
 
 	// wait until we successfully create a raft group or join an existing one
 	raftManager := raftmanager.New(nodeID, dispatcher, peerDB, hconfig)
-	joinCtx, joinCancel := context.WithTimeout(ctx, time.Minute)
+	joinCtx, joinCancel := context.WithTimeout(ctx, hconfig.Raft.InitialJoinDuration)
 	defer joinCancel()
 	if err := raftManager.CreateOrJoin(joinCtx); err != nil {
 		return fmt.Errorf("failure to join raft : %v", err)