From 049c26229db06d84b604db2c6ada2d52da2226ff Mon Sep 17 00:00:00 2001
From: Alexander Matyushentsev <AMatyushentsev@gmail.com>
Date: Wed, 3 May 2023 04:31:29 -0700
Subject: [PATCH] Let agent retry to connecting to server (#1728)

Closes https://github.com/woodpecker-ci/woodpecker/issues/1721

PR introduces connection retries during agent startup.

Signed-off-by: Alexander Matyushentsev <AMatyushentsev@gmail.com>
---
 cmd/agent/agent.go                            | 19 ++++++++++++++++++-
 cmd/agent/flags.go                            | 12 ++++++++++++
 cmd/agent/main.go                             |  2 +-
 .../docs/30-administration/15-agent-config.md | 18 ++++++++++++++++++
 4 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/cmd/agent/agent.go b/cmd/agent/agent.go
index 8c285f2aa..e57447d25 100644
--- a/cmd/agent/agent.go
+++ b/cmd/agent/agent.go
@@ -32,10 +32,12 @@ import (
 	"github.com/tevino/abool"
 	"github.com/urfave/cli/v2"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 	grpccredentials "google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/keepalive"
 	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
 
 	"github.com/woodpecker-ci/woodpecker/agent"
 	agentRpc "github.com/woodpecker-ci/woodpecker/agent/rpc"
@@ -46,7 +48,7 @@ import (
 	"github.com/woodpecker-ci/woodpecker/version"
 )
 
-func loop(c *cli.Context) error {
+func run(c *cli.Context) error {
 	hostname := c.String("hostname")
 	if len(hostname) == 0 {
 		hostname, _ = os.Hostname()
@@ -246,6 +248,21 @@ func loop(c *cli.Context) error {
 	return nil
 }
 
+func runWithRetry(context *cli.Context) error {
+	retryCount := context.Int("connect-retry-count")
+	retryDelay := context.Duration("connect-retry-delay")
+	var err error
+	for i := 0; i < retryCount; i++ {
+		if err = run(context); status.Code(err) == codes.Unavailable {
+			log.Warn().Err(err).Msg(fmt.Sprintf("cannot connect to server, retrying in %v", retryDelay))
+			time.Sleep(retryDelay)
+		} else {
+			break
+		}
+	}
+	return err
+}
+
 func stringSliceAddToMap(sl []string, m map[string]string) error {
 	if m == nil {
 		m = make(map[string]string)
diff --git a/cmd/agent/flags.go b/cmd/agent/flags.go
index ba40bd881..8a191b791 100644
--- a/cmd/agent/flags.go
+++ b/cmd/agent/flags.go
@@ -190,4 +190,16 @@ var flags = []cli.Flag{
 		Usage:   "backend k8s additional worker pod annotations",
 		Value:   "",
 	},
+	&cli.IntFlag{
+		EnvVars: []string{"WOODPECKER_CONNECT_RETRY_COUNT"},
+		Name:    "connect-retry-count",
+		Usage:   "number of times to retry connecting to the server",
+		Value:   5,
+	},
+	&cli.DurationFlag{
+		EnvVars: []string{"WOODPECKER_CONNECT_RETRY_DELAY"},
+		Name:    "connect-retry-delay",
+		Usage:   "duration to wait before retrying to connect to the server",
+		Value:   time.Second * 2,
+	},
 }
diff --git a/cmd/agent/main.go b/cmd/agent/main.go
index c6fa71d07..7543f87fb 100644
--- a/cmd/agent/main.go
+++ b/cmd/agent/main.go
@@ -29,7 +29,7 @@ func main() {
 	app.Name = "woodpecker-agent"
 	app.Version = version.String()
 	app.Usage = "woodpecker agent"
-	app.Action = loop
+	app.Action = runWithRetry
 	app.Commands = []*cli.Command{
 		{
 			Name:   "ping",
diff --git a/docs/docs/30-administration/15-agent-config.md b/docs/docs/30-administration/15-agent-config.md
index 4c31af700..ac7e15c55 100644
--- a/docs/docs/30-administration/15-agent-config.md
+++ b/docs/docs/30-administration/15-agent-config.md
@@ -132,3 +132,21 @@ See [Docker backend configuration](backends/docker/#configuration)
 ### `WOODPECKER_BACKEND_SSH_*`
 
 See [SSH backend configuration](backends/ssh/#configuration)
+
+## Advanced Settings
+
+:::warning
+Only change these If you know what you do.
+:::
+
+### `WOODPECKER_CONNECT_RETRY_COUNT`
+
+> Default: `5`
+
+Configures number of times agent retries to connect to the server.
+
+### `WOODPECKER_CONNECT_RETRY_DELAY`
+
+> Default: `2s`
+
+Configures delay between agent connection retries to the server.