Blockade testing

Fix an issue during retrieveSnapshot, where the membership updates from
Zero haven't yet rolled in, with info about Alpha leader. Instead of
relying solely upon that, we use the RaftContext in Snapshot to
determine the leader, and get the connection.

Do not destroy blockade, in case the tests fail. We'd like to be able to
get the logs from the Dgraph servers.

Author: manishrjain

contrib/blockade/main.go

@@ -131,12 +131,6 @@ func waitForHealthy() error {
 }
 
 func runTests() error {
-	defer func() {
-		if err := run(ctxb, "blockade destroy"); err != nil {
-			log.Fatalf("While destroying: %v", err)
-		}
-	}()
-
 	for {
 		if err := waitForHealthy(); err != nil {
 			fmt.Printf("Error while waitForHealthy: %v\n.", err)
@@ -181,6 +175,15 @@ func main() {
 	if err := run(ctxb, "blockade up"); err != nil {
 		log.Fatal(err)
 	}
+	// This defer can be moved within runTests, if we want to destroy blockade,
+	// in case our tests fail. We don't want to do that, because then we won't
+	// be able to get the logs.
+	defer func() {
+		if err := run(ctxb, "blockade destroy"); err != nil {
+			log.Fatalf("While destroying: %v", err)
+		}
+	}()
+
 	if err := runTests(); err != nil {
 		os.Exit(1)
 	}

worker/draft.go

@@ -531,9 +531,29 @@ func (n *node) Snapshot() (*pb.Snapshot, error) {
 }
 
 func (n *node) retrieveSnapshot(snap pb.Snapshot) error {
-	pool, err := n.leaderBlocking()
-	if err != nil {
-		return err
+	// In some edge cases, the Zero leader might not have been able to update
+	// the status of Alpha leader. So, instead of blocking forever on waiting
+	// for Zero to send us the updates info about the leader, we can just use
+	// the Snapshot RaftContext, which contains the address of the leader.
+	var pool *conn.Pool
+	addr := snap.Context.GetAddr()
+	glog.V(2).Infof("Snapshot.RaftContext.Addr: %q", addr)
+	if len(addr) > 0 {
+		p, err := conn.Get().Get(addr)
+		if err != nil {
+			glog.V(2).Infof("conn.Get(%q) Error: %v", addr, err)
+		} else {
+			pool = p
+			glog.V(2).Infof("Leader connection picked from RaftContext")
+		}
+	}
+	if pool == nil {
+		glog.V(2).Infof("No leader conn from RaftContext. Using membership state.")
+		p, err := n.leaderBlocking()
+		if err != nil {
+			return err
+		}
+		pool = p
 	}
 
 	// Need to clear pl's stored in memory for the case when retrieving snapshot with