From 31d22f711ef5974f5d561061063f794927ec2a7f Mon Sep 17 00:00:00 2001
From: Geoff Bourne <itzgeoff@gmail.com>
Date: Sun, 28 Jun 2020 15:52:14 -0500
Subject: [PATCH] Added support for retries on java and bedrock status

---
 bedrock_status.go | 33 +++++++++++++++++++++---------
 java_status.go    | 51 ++++++++++++++++++++++++++++++++---------------
 2 files changed, 59 insertions(+), 25 deletions(-)

diff --git a/bedrock_status.go b/bedrock_status.go
index 02de7e6..3c9ac58 100644
--- a/bedrock_status.go
+++ b/bedrock_status.go
@@ -9,11 +9,15 @@ import (
 	"log"
 	"net"
 	"strconv"
+	"time"
 )
 
 type statusBedrockCmd struct {
 	Host string `default:"localhost"`
 	Port int    `default:"19132"`
+
+	RetryInterval time.Duration `usage:"if retry-limit is non-zero, status will be retried at this interval" default:"10s"`
+	RetryLimit    int           `usage:"if non-zero, failed status will be retried this many times before exiting"`
 }
 
 func (c *statusBedrockCmd) Name() string {
@@ -38,16 +42,27 @@ func (c *statusBedrockCmd) SetFlags(flags *flag.FlagSet) {
 
 func (c *statusBedrockCmd) Execute(ctx context.Context, f *flag.FlagSet, args ...interface{}) subcommands.ExitStatus {
 	address := net.JoinHostPort(c.Host, strconv.Itoa(c.Port))
-
-	info, err := PingBedrockServer(address)
-	if err != nil {
-		log.Fatal(err)
-		return subcommands.ExitFailure
+	if c.RetryInterval <= 0 {
+		c.RetryInterval = 1 * time.Second
 	}
 
-	fmt.Printf("%s : version=%s online=%d max=%d",
-		address,
-		info.Version, info.Players, info.MaxPlayers)
+	for {
+		info, err := PingBedrockServer(address)
+		if err != nil {
+			if c.RetryLimit > 0 {
+				c.RetryLimit--
+				time.Sleep(c.RetryInterval)
+				continue
+			}
+			log.Fatal(err)
+			return subcommands.ExitFailure
+		}
+
+		fmt.Printf("%s : version=%s online=%d max=%d",
+			address,
+			info.Version, info.Players, info.MaxPlayers)
+
+		return subcommands.ExitSuccess
+	}
 
-	return subcommands.ExitSuccess
 }
diff --git a/java_status.go b/java_status.go
index 2a69bab..3f0ca2e 100644
--- a/java_status.go
+++ b/java_status.go
@@ -9,11 +9,15 @@ import (
 	"github.com/itzg/go-flagsfiller"
 	"log"
 	"os"
+	"time"
 )
 
 type statusCmd struct {
 	Host string `default:"localhost" usage:"hostname of the Minecraft server" env:"MC_HOST"`
 	Port int    `default:"25565" usage:"port of the Minecraft server" env:"MC_PORT"`
+
+	RetryInterval time.Duration `usage:"if retry-limit is non-zero, status will be retried at this interval" default:"10s"`
+	RetryLimit    int           `usage:"if non-zero, failed status will be retried this many times before exiting"`
 }
 
 func (c *statusCmd) Name() string {
@@ -38,24 +42,39 @@ func (c *statusCmd) SetFlags(flags *flag.FlagSet) {
 
 func (c *statusCmd) Execute(ctx context.Context, f *flag.FlagSet, args ...interface{}) subcommands.ExitStatus {
 	pinger := mcpinger.New(c.Host, uint16(c.Port))
-
-	info, err := pinger.Ping()
-	if err != nil {
-		_, _ = fmt.Fprintf(os.Stderr, "failed to ping %s:%d : %s", c.Host, c.Port, err)
-		return subcommands.ExitFailure
+	if c.RetryInterval <= 0 {
+		c.RetryInterval = 1 * time.Second
 	}
 
-	// While server is starting up it will answer pings, but respond with empty JSON object.
-	// As such, we'll sanity check the max players value to see if a zero-value has been
-	// provided for info.
-	if info.Players.Max == 0 {
-		_, _ = fmt.Fprintf(os.Stderr, "server not ready %s:%d", c.Host, c.Port)
-		return subcommands.ExitFailure
-	}
+	for {
+		info, err := pinger.Ping()
+		if err != nil {
+			if c.RetryLimit > 0 {
+				c.RetryLimit--
+				time.Sleep(c.RetryInterval)
+				continue
+			}
+			_, _ = fmt.Fprintf(os.Stderr, "failed to ping %s:%d : %s", c.Host, c.Port, err)
+			return subcommands.ExitFailure
+		}
+
+		// While server is starting up it will answer pings, but respond with empty JSON object.
+		// As such, we'll sanity check the max players value to see if a zero-value has been
+		// provided for info.
+		if info.Players.Max == 0 {
+			if c.RetryLimit > 0 {
+				c.RetryLimit--
+				time.Sleep(c.RetryInterval)
+				continue
+			}
+			_, _ = fmt.Fprintf(os.Stderr, "server not ready %s:%d", c.Host, c.Port)
+			return subcommands.ExitFailure
+		}
 
-	fmt.Printf("%s:%d : version=%s online=%d max=%d motd='%s'",
-		c.Host, c.Port,
-		info.Version.Name, info.Players.Online, info.Players.Max, info.Description.Text)
+		fmt.Printf("%s:%d : version=%s online=%d max=%d motd='%s'",
+			c.Host, c.Port,
+			info.Version.Name, info.Players.Online, info.Players.Max, info.Description.Text)
 
-	return subcommands.ExitSuccess
+		return subcommands.ExitSuccess
+	}
 }