From 3ea795a1ad109c6891c3a39a1030688e7a6c8847 Mon Sep 17 00:00:00 2001
From: Nikita Kryuchkov <nkryuchkov10@gmail.com>
Date: Thu, 23 Jul 2020 17:26:03 +0300
Subject: [PATCH] Fix health endpoint timeout

---
 pkg/hypervisor/hypervisor.go       |  5 +-
 pkg/routefinder/rfclient/client.go |  9 +++-
 pkg/visor/rpc.go                   | 83 +++++++++++++++++++-----------
 3 files changed, 62 insertions(+), 35 deletions(-)

diff --git a/pkg/hypervisor/hypervisor.go b/pkg/hypervisor/hypervisor.go
index 2e6eb79bd..51d415427 100644
--- a/pkg/hypervisor/hypervisor.go
+++ b/pkg/hypervisor/hypervisor.go
@@ -34,8 +34,7 @@ import (
 )
 
 const (
-	healthTimeout = 5 * time.Second
-	httpTimeout   = 30 * time.Second
+	httpTimeout = 30 * time.Second
 )
 
 const (
@@ -366,7 +365,7 @@ func (hv *Hypervisor) getHealth() http.HandlerFunc {
 		}
 
 		resCh := make(chan healthRes)
-		tCh := time.After(healthTimeout)
+		tCh := time.After(visor.HealthTimeout)
 
 		go func() {
 			hi, err := ctx.RPC.Health()
diff --git a/pkg/routefinder/rfclient/client.go b/pkg/routefinder/rfclient/client.go
index c5c5913b3..62ae8e1f3 100644
--- a/pkg/routefinder/rfclient/client.go
+++ b/pkg/routefinder/rfclient/client.go
@@ -128,8 +128,13 @@ func (c *apiClient) FindRoutes(ctx context.Context, rts []routing.PathEdges, opt
 }
 
 // Health checks route finder health.
-func (c *apiClient) Health(_ context.Context) (int, error) {
-	res, err := http.Get(c.addr + "/health")
+func (c *apiClient) Health(ctx context.Context) (int, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.addr+"/health", nil)
+	if err != nil {
+		return 0, err
+	}
+
+	res, err := http.DefaultClient.Do(req)
 	if err != nil {
 		return 0, err
 	}
diff --git a/pkg/visor/rpc.go b/pkg/visor/rpc.go
index 5b8e84438..e26f4daca 100644
--- a/pkg/visor/rpc.go
+++ b/pkg/visor/rpc.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/rpc"
+	"sync"
 	"time"
 
 	"github.com/SkycoinProject/dmsg/buildinfo"
@@ -23,7 +24,8 @@ import (
 
 const (
 	// RPCPrefix is the prefix used with all RPC calls.
-	RPCPrefix = "app-visor"
+	RPCPrefix     = "app-visor"
+	HealthTimeout = 5 * time.Second
 )
 
 var (
@@ -74,7 +76,8 @@ type HealthInfo struct {
 func (r *RPC) Health(_ *struct{}, out *HealthInfo) (err error) {
 	defer rpcutil.LogCall(r.log, "Health", nil)(out, &err)
 
-	ctx := context.Background()
+	ctx, cancel := context.WithTimeout(context.Background(), HealthTimeout/2)
+	defer cancel()
 
 	out.TransportDiscovery = http.StatusNotFound
 	out.RouteFinder = http.StatusNotFound
@@ -82,26 +85,36 @@ func (r *RPC) Health(_ *struct{}, out *HealthInfo) (err error) {
 	out.UptimeTracker = http.StatusNotFound
 	out.AddressResolver = http.StatusNotFound
 
-	if tdClient := r.visor.TpDiscClient(); tdClient != nil {
-		tdStatus, err := tdClient.Health(ctx)
-		if err != nil {
-			r.log.WithError(err).Warnf("Failed to check transport discovery health")
+	var wg sync.WaitGroup
 
-			out.TransportDiscovery = http.StatusInternalServerError
-		}
+	if tdClient := r.visor.TpDiscClient(); tdClient != nil {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			tdStatus, err := tdClient.Health(ctx)
+			if err != nil {
+				r.log.WithError(err).Warnf("Failed to check transport discovery health")
+
+				out.TransportDiscovery = http.StatusInternalServerError
+			}
 
-		out.TransportDiscovery = tdStatus
+			out.TransportDiscovery = tdStatus
+		}()
 	}
 
 	if rfClient := r.visor.RouteFinderClient(); rfClient != nil {
-		rfStatus, err := rfClient.Health(ctx)
-		if err != nil {
-			r.log.WithError(err).Warnf("Failed to check route finder health")
-
-			out.RouteFinder = http.StatusInternalServerError
-		}
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			rfStatus, err := rfClient.Health(ctx)
+			if err != nil {
+				r.log.WithError(err).Warnf("Failed to check route finder health")
+
+				out.RouteFinder = http.StatusInternalServerError
+			}
 
-		out.RouteFinder = rfStatus
+			out.RouteFinder = rfStatus
+		}()
 	}
 
 	// TODO(evanlinjin): This should actually poll the setup nodes services.
@@ -112,27 +125,37 @@ func (r *RPC) Health(_ *struct{}, out *HealthInfo) (err error) {
 	}
 
 	if utClient := r.visor.UptimeTrackerClient(); utClient != nil {
-		utStatus, err := utClient.Health(ctx)
-		if err != nil {
-			r.log.WithError(err).Warnf("Failed to check uptime tracker health")
-
-			out.UptimeTracker = http.StatusInternalServerError
-		}
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			utStatus, err := utClient.Health(ctx)
+			if err != nil {
+				r.log.WithError(err).Warnf("Failed to check uptime tracker health")
+
+				out.UptimeTracker = http.StatusInternalServerError
+			}
 
-		out.UptimeTracker = utStatus
+			out.UptimeTracker = utStatus
+		}()
 	}
 
 	if arClient := r.visor.AddressResolverClient(); arClient != nil {
-		arStatus, err := arClient.Health(ctx)
-		if err != nil {
-			r.log.WithError(err).Warnf("Failed to check address resolver health")
-
-			out.AddressResolver = http.StatusInternalServerError
-		}
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			arStatus, err := arClient.Health(ctx)
+			if err != nil {
+				r.log.WithError(err).Warnf("Failed to check address resolver health")
+
+				out.AddressResolver = http.StatusInternalServerError
+			}
 
-		out.AddressResolver = arStatus
+			out.AddressResolver = arStatus
+		}()
 	}
 
+	wg.Wait()
+
 	return nil
 }