From 42be4bdc0ffa31865d27da3547667ba408ea59e0 Mon Sep 17 00:00:00 2001
From: Cameron Ayer <camayer92@gmail.com>
Date: Thu, 2 Apr 2020 16:44:51 -0400
Subject: [PATCH] satellite/contact: add timeout to PingBack method

Change-Id: I2ec2f82e2e10d8be16f82e9de13ce42358e47c98
---
 private/testplanet/satellite.go             |  3 +++
 satellite/api.go                            |  2 +-
 satellite/contact/service.go                | 15 +++++++++++++--
 satellite/core.go                           |  2 +-
 scripts/testdata/satellite-config.yaml.lock |  3 +++
 5 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/private/testplanet/satellite.go b/private/testplanet/satellite.go
index 1b1be2a40..accba96a1 100644
--- a/private/testplanet/satellite.go
+++ b/private/testplanet/satellite.go
@@ -310,6 +310,9 @@ func (planet *Planet) newSatellites(count int, satelliteDatabases satellitedbtes
 			Admin: admin.Config{
 				Address: "127.0.0.1:0",
 			},
+			Contact: contact.Config{
+				Timeout: 1 * time.Minute,
+			},
 			Overlay: overlay.Config{
 				Node: overlay.NodeSelectionConfig{
 					UptimeCount:      0,
diff --git a/satellite/api.go b/satellite/api.go
index 2790abfbd..73b78e5f7 100644
--- a/satellite/api.go
+++ b/satellite/api.go
@@ -289,7 +289,7 @@ func NewAPI(log *zap.Logger, full *identity.FullIdentity, db DB,
 			Type:    pb.NodeType_SATELLITE,
 			Version: *pbVersion,
 		}
-		peer.Contact.Service = contact.NewService(peer.Log.Named("contact:service"), self, peer.Overlay.Service, peer.DB.PeerIdentities(), peer.Dialer)
+		peer.Contact.Service = contact.NewService(peer.Log.Named("contact:service"), self, peer.Overlay.Service, peer.DB.PeerIdentities(), peer.Dialer, config.Contact.Timeout)
 		peer.Contact.Endpoint = contact.NewEndpoint(peer.Log.Named("contact:endpoint"), peer.Contact.Service)
 		pbgrpc.RegisterNodeServer(peer.Server.GRPC(), peer.Contact.Endpoint)
 		if err := pb.DRPCRegisterNode(peer.Server.DRPC(), peer.Contact.Endpoint); err != nil {
diff --git a/satellite/contact/service.go b/satellite/contact/service.go
index abb5ee510..7092cb938 100644
--- a/satellite/contact/service.go
+++ b/satellite/contact/service.go
@@ -7,6 +7,7 @@ import (
 	"context"
 	"fmt"
 	"sync"
+	"time"
 
 	"github.com/zeebo/errs"
 	"go.uber.org/zap"
@@ -20,7 +21,8 @@ import (
 
 // Config contains configurable values for contact service
 type Config struct {
-	ExternalAddress string `user:"true" help:"the public address of the node, useful for nodes behind NAT" default:""`
+	ExternalAddress string        `user:"true" help:"the public address of the node, useful for nodes behind NAT" default:""`
+	Timeout         time.Duration `help:"timeout for pinging storage nodes" default:"10m0s"`
 }
 
 // Service is the contact service between storage nodes and satellites.
@@ -37,16 +39,19 @@ type Service struct {
 	overlay *overlay.Service
 	peerIDs overlay.PeerIdentities
 	dialer  rpc.Dialer
+
+	timeout time.Duration
 }
 
 // NewService creates a new contact service.
-func NewService(log *zap.Logger, self *overlay.NodeDossier, overlay *overlay.Service, peerIDs overlay.PeerIdentities, dialer rpc.Dialer) *Service {
+func NewService(log *zap.Logger, self *overlay.NodeDossier, overlay *overlay.Service, peerIDs overlay.PeerIdentities, dialer rpc.Dialer, timeout time.Duration) *Service {
 	return &Service{
 		log:     log,
 		self:    self,
 		overlay: overlay,
 		peerIDs: peerIDs,
 		dialer:  dialer,
+		timeout: timeout,
 	}
 }
 
@@ -64,6 +69,12 @@ func (service *Service) Close() error { return nil }
 func (service *Service) PingBack(ctx context.Context, address string, peerID storj.NodeID) (_ bool, _ string, err error) {
 	defer mon.Task()(&ctx)(&err)
 
+	if service.timeout > 0 {
+		var cancel func()
+		ctx, cancel = context.WithTimeout(ctx, service.timeout)
+		defer cancel()
+	}
+
 	pingNodeSuccess := true
 	var pingErrorMessage string
 
diff --git a/satellite/core.go b/satellite/core.go
index 5ddec3eeb..2ee4eefc2 100644
--- a/satellite/core.go
+++ b/satellite/core.go
@@ -218,7 +218,7 @@ func New(log *zap.Logger, full *identity.FullIdentity, db DB,
 			Type:    pb.NodeType_SATELLITE,
 			Version: *pbVersion,
 		}
-		peer.Contact.Service = contact.NewService(peer.Log.Named("contact:service"), self, peer.Overlay.Service, peer.DB.PeerIdentities(), peer.Dialer)
+		peer.Contact.Service = contact.NewService(peer.Log.Named("contact:service"), self, peer.Overlay.Service, peer.DB.PeerIdentities(), peer.Dialer, config.Contact.Timeout)
 		peer.Services.Add(lifecycle.Item{
 			Name:  "contact:service",
 			Close: peer.Contact.Service.Close,
diff --git a/scripts/testdata/satellite-config.yaml.lock b/scripts/testdata/satellite-config.yaml.lock
index 247e87cd2..7d517d8d7 100644
--- a/scripts/testdata/satellite-config.yaml.lock
+++ b/scripts/testdata/satellite-config.yaml.lock
@@ -106,6 +106,9 @@ compensation.withheld-percents: 75,75,75,50,50,50,25,25,25,0,0,0,0,0,0
 # the public address of the node, useful for nodes behind NAT
 contact.external-address: ""
 
+# timeout for pinging storage nodes
+# contact.timeout: 10m0s
+
 # satellite database connection string
 # database: postgres://