From cb2e36886580301b0c0493d91c554141ba9bcf57 Mon Sep 17 00:00:00 2001
From: Mark Hindess <mark.hindess@gmail.com>
Date: Sun, 23 Jul 2023 15:25:09 +0100
Subject: [PATCH] fix: add retry logic to ListPartitionReassignments

Signed-off-by: Mark Hindess <mark.hindess@gmail.com>
Signed-off-by: Ioan Zicu <ioan.zicu@nokia.com>
---
 admin.go | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/admin.go b/admin.go
index ee4cd4ca68..29eeca1c61 100644
--- a/admin.go
+++ b/admin.go
@@ -556,13 +556,20 @@ func (ca *clusterAdmin) ListPartitionReassignments(topic string, partitions []in
 
 	request.AddBlock(topic, partitions)
 
-	b, err := ca.Controller()
-	if err != nil {
-		return nil, err
-	}
-	_ = b.Open(ca.client.Config())
+	var rsp *ListPartitionReassignmentsResponse
+	err = ca.retryOnError(isErrNoController, func() error {
+		b, err := ca.Controller()
+		if err != nil {
+			return err
+		}
+		_ = b.Open(ca.client.Config())
 
-	rsp, err := b.ListPartitionReassignments(request)
+		rsp, err = b.ListPartitionReassignments(request)
+		if isErrNoController(err) {
+			_, _ = ca.refreshController()
+		}
+		return err
+	})
 
 	if err == nil && rsp != nil {
 		return rsp.TopicStatus, nil