feat(wait): define wait pkg and changes to enhance pkg experience

.gitignore

@@ -1,2 +1,3 @@
 bin/
-./wait-for-it
+*wait-for-it
+.idea/

README.md

@@ -50,7 +50,10 @@ This package is adapted from [vishnubob/wait-for-it](https://github.com/vishnubo
 
 Since [vishnubob/wait-for-it](https://github.com/vishnubob/wait-for-it) is a bash script, it does not work directly with minimal containers like [scratch](https://hub.docker.com/_/scratch), which are commonly used to run binaries.
 
-With the help of this package, you can generate a binary, which can run inside minimal Docker containers and wait for a TCP connection such as a `mysql` database. You can find an example here: [csivitu/bl0b](https://github.com/csivitu/bl0b/blob/master/docker-compose.yml).
+With the help of this package, you can:
+1. Generate a binary, which can run inside minimal Docker containers and wait for a TCP connection such as a `mysql` database. You can find an example here: [csivitu/bl0b](https://github.com/csivitu/bl0b/blob/master/docker-compose.yml).
+2. Consume the `wait` package internally within your Go microservices.
+This would allow your services running natively to have the same wait setup as that of when running within a Docker environment.
 
 
 ### Built With
@@ -107,12 +110,14 @@ Use `wait-for-it -h` to display the following list.
 
 ```
 Usage of wait-for-it:
+  -m int
+        Max service timeout to retry request in seconds, zero for no max service timeout (default 30)
   -q    Quiet, don't output any status messages
   -s    Only execute subcommand if the test succeeds
   -t int
-        Timeout in seconds, zero for no timeout (default 15)
+        Service request timeout in seconds, zero for no timeout (default 15)
   -w host:port
-        Services to be waiting for, in the form host:port
+        Dependency services to be waiting for, in the form host:port
 ```
 
 You can run any executable after passing ` -- `, like in the examples below.
@@ -129,8 +134,8 @@ wait-for-it -w google.com:80 -w localhost:27017 -t 30 -- echo "Waiting for 30 se
 
 ```sh
 $ wait-for-it -w abcd:80 -s -t 5 -- echo "Done\!"
-wait-for-it: waiting 5 seconds for abcd:80
-wait-for-it: timeout occured after waiting for 5 seconds
+wait-for-it: waiting 5 seconds for abcd:80 for a max of 10 seconds
+wait-for-it: failed to dial service abcd:80 with err: dial tcp: lookup abcd on 172.24.128.1:53: no such host
 wait-for-it: strict mode, refusing to execute subprocess
 ```
 

main.go

@@ -3,46 +3,50 @@ package main
 import (
 	"flag"
 	"fmt"
+	"github.com/roerohan/wait-for-it/pkg/wait"
 	"os"
+	"time"
 )
 
 var (
-	timeout  int
-	services Services
-	quiet    bool
-	strict   bool
+	reqTimeout int
+	maxTimeout int
+	services   wait.Services
+	quiet      bool
+	strict     bool
 )
 
 func init() {
-	flag.IntVar(&timeout, "t", 15, "Timeout in seconds, zero for no timeout")
+	flag.IntVar(&reqTimeout, "t", 15, "Service request timeout in seconds, zero for no timeout")
+	flag.IntVar(&maxTimeout, "m", 30, "Max service timeout to retry request in seconds, zero for no max service timeout")
 	flag.BoolVar(&quiet, "q", false, "Quiet, don't output any status messages")
 	flag.BoolVar(&strict, "s", false, "Only execute subcommand if the test succeeds")
-	flag.Var(&services, "w", "Services to be waiting for, in the form `host:port`")
+	flag.Var(&services, "w", "Dependency services to be waiting for, in the form `host:port`")
 }
 
 // Log is used to log with prefix wait-for-it:
-func Log(message string) {
+func log(message string) {
 	if quiet {
 		return
 	}
 
-	fmt.Println("wait-for-it: " + message)
+	wait.Log(message)
 }
 
 func main() {
 	flag.Parse()
 	args := os.Args
 
 	if len(services) != 0 {
-		Log(fmt.Sprintf("waiting %d seconds for %s", timeout, services.String()))
-		ok := services.Wait(timeout)
-
-		if !ok {
-			Log(fmt.Sprintf("timeout occured after waiting for %d seconds", timeout))
-			if strict {
-				Log("strict mode, refusing to execute subprocess")
-				os.Exit(1)
-			}
+		log(fmt.Sprintf("waiting %d seconds for %s for a max of %d seconds", reqTimeout, services.String(), maxTimeout))
+		err := wait.ForDependencies(services, time.Duration(reqTimeout), time.Duration(maxTimeout))
+		if err != nil {
+			log(fmt.Sprintf("wait.ForDependencies failed with err %v", err))
+			os.Exit(1)
+		}
+		if strict {
+			log("strict mode, refusing to execute subprocess")
+			os.Exit(1)
 		}
 	}
 

pkg/wait/log.go

@@ -0,0 +1,8 @@
+package wait
+
+import "fmt"
+
+// Log is used to log with prefix wait-for-it
+func Log(message string) {
+	fmt.Println("wait-for-it: " + message)
+}

pkg/wait/wait.go

@@ -0,0 +1,134 @@
+package wait
+
+import (
+	"errors"
+	"fmt"
+	"net"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+)
+
+// Services is a string array storing the services that are to be waited for
+type Services []Service
+
+// Service is a struct meant to denote a service hostname:port with a wait condition upon start up
+type Service struct {
+	hostname string
+	port     int
+}
+
+var (
+	// ErrServiceMaxTimeout is the error message to use in case a max service startup timeout is exceeded.
+	ErrServiceMaxTimeout = fmt.Errorf("max service startup timeout duration exceeded waiting for service dependencies")
+)
+
+// String prints out the human-readable Service hostname:port string.
+func (s *Service) String() string {
+	return fmt.Sprintf("%s:%d", s.hostname, s.port)
+}
+
+// NewService creates a wait.Service type.
+func NewService(hostname string, port int) Service {
+	return Service{
+		hostname: hostname,
+		port:     port,
+	}
+}
+
+// Set is used to append a Service to the slice of Services,
+// to implement the interface flag.Value
+func (s *Services) Set(value string) error {
+	const separator = ":"
+
+	// Note: serviceInfo[0] = hostname, serviceInfo[1] = port
+	serviceInfo := strings.Split(value, separator)
+	port, err := strconv.Atoi(serviceInfo[1])
+	if err != nil {
+		return err
+	}
+
+	service := NewService(serviceInfo[0], port)
+	*s = append(*s, service)
+	return nil
+}
+
+// String returns a string representation of the flag,
+// to implement the interface flag.Value
+func (s *Services) String() string {
+	var sb strings.Builder
+	const formatter string = ", "
+	for _, service := range *s {
+		sb.WriteString(service.String())
+		sb.WriteString(formatter)
+	}
+
+	// trim the last comma that was added for last service
+	return strings.TrimSuffix(sb.String(), formatter)
+}
+
+// ForDependencies allows the service to wait for its dependencies to be up and ready for a configurable amount of time.
+// If the service dependency request timeout is reached and the dependent services are not yet available,
+// then the timeout wait interval will continue until the dependencies are up for a maximum wait time of maxTimeout.
+func ForDependencies(waitServices Services, serviceRequestTimeout, maxTimeout time.Duration) error {
+	if len(waitServices) == 0 {
+		return nil
+	}
+
+	success := make(chan bool, 1)
+	ok := wait(waitServices, serviceRequestTimeout)
+	if ok {
+		success <- true
+	}
+
+	// return err if service wait time exceeds ServiceMaxTimeout time
+	select {
+	case <-success:
+		return nil
+	case <-time.After(maxTimeout * time.Second):
+		return ErrServiceMaxTimeout
+	}
+}
+
+func wait(waitServices Services, waitTimeOut time.Duration) bool {
+	now := time.Now()
+
+	var wg sync.WaitGroup
+	wg.Add(len(waitServices))
+
+	success := make(chan bool, 1)
+
+	go func() {
+		for _, service := range waitServices {
+			go waitOne(service, &wg, now)
+		}
+		wg.Wait()
+		success <- true
+	}()
+
+	select {
+	case <-success:
+		return true
+	case <-time.After(waitTimeOut * time.Second):
+		return false
+	}
+
+}
+
+func waitOne(service Service, wg *sync.WaitGroup, start time.Time) {
+	defer wg.Done()
+	for {
+		_, err := net.Dial("tcp", service.String())
+		if err == nil {
+			Log(fmt.Sprintf("%s is available after %s", service.String(), time.Since(start)))
+			break
+		}
+		opErr, ok := err.(*net.OpError)
+		if ok && errors.Is(err, opErr) {
+			Log(fmt.Sprintf("failed to dial service %s with err: %s", service.String(), opErr.Error()))
+			break
+		}
+		time.Sleep(time.Second)
+	}
+}