From cbe083067986001df253dafe3c9e37c9e811d8a6 Mon Sep 17 00:00:00 2001
From: Andrew Gillis <gammazero@users.noreply.github.com>
Date: Mon, 31 Jul 2023 11:07:51 -0700
Subject: [PATCH] Option to set distance update timeout (#48)

---
 pkg/dtrack/distance_tracker.go | 20 +++++++++++++-------
 pkg/provider/provider.go       | 21 ++++++++++++++++++---
 version.json                   |  2 +-
 3 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/pkg/dtrack/distance_tracker.go b/pkg/dtrack/distance_tracker.go
index cec8dbd..2a3f96d 100644
--- a/pkg/dtrack/distance_tracker.go
+++ b/pkg/dtrack/distance_tracker.go
@@ -34,14 +34,14 @@ type distTrack struct {
 	errType int
 }
 
-func RunDistanceTracker(ctx context.Context, include, exclude map[peer.ID]struct{}, provCache *pcache.ProviderCache, depthLimit int64, updateIn time.Duration) <-chan DistanceUpdate {
+func RunDistanceTracker(ctx context.Context, include, exclude map[peer.ID]struct{}, provCache *pcache.ProviderCache, depthLimit int64, updateIn, timeout time.Duration) <-chan DistanceUpdate {
 	updates := make(chan DistanceUpdate)
-	go runTracker(ctx, include, exclude, provCache, updateIn, depthLimit, updates)
+	go runTracker(ctx, include, exclude, provCache, updateIn, timeout, depthLimit, updates)
 
 	return updates
 }
 
-func runTracker(ctx context.Context, include, exclude map[peer.ID]struct{}, provCache *pcache.ProviderCache, updateIn time.Duration, depthLimit int64, updates chan<- DistanceUpdate) {
+func runTracker(ctx context.Context, include, exclude map[peer.ID]struct{}, provCache *pcache.ProviderCache, updateIn, timeout time.Duration, depthLimit int64, updates chan<- DistanceUpdate) {
 	defer close(updates)
 
 	var lookForNew bool
@@ -78,7 +78,7 @@ func runTracker(ctx context.Context, include, exclude map[peer.ID]struct{}, prov
 					}
 				}
 			}
-			updateTracks(ctx, provCache, tracks, depthLimit, updates)
+			updateTracks(ctx, provCache, tracks, timeout, depthLimit, updates)
 			timer.Reset(updateIn)
 		case <-ctx.Done():
 			return
@@ -86,13 +86,19 @@ func runTracker(ctx context.Context, include, exclude map[peer.ID]struct{}, prov
 	}
 }
 
-func updateTracks(ctx context.Context, provCache *pcache.ProviderCache, tracks map[peer.ID]*distTrack, depthLimit int64, updates chan<- DistanceUpdate) {
+func updateTracks(ctx context.Context, provCache *pcache.ProviderCache, tracks map[peer.ID]*distTrack, timeout time.Duration, depthLimit int64, updates chan<- DistanceUpdate) {
 	for providerID, track := range tracks {
-		updateTrack(ctx, providerID, track, provCache, depthLimit, updates)
+		updateTrack(ctx, providerID, track, provCache, timeout, depthLimit, updates)
 	}
 }
 
-func updateTrack(ctx context.Context, pid peer.ID, track *distTrack, provCache *pcache.ProviderCache, depthLimit int64, updates chan<- DistanceUpdate) {
+func updateTrack(ctx context.Context, pid peer.ID, track *distTrack, provCache *pcache.ProviderCache, timeout time.Duration, depthLimit int64, updates chan<- DistanceUpdate) {
+	if timeout != 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, timeout)
+		defer cancel()
+	}
+
 	pinfo, err := provCache.Get(ctx, pid)
 	if err != nil {
 		return
diff --git a/pkg/provider/provider.go b/pkg/provider/provider.go
index 27ce9dd..9c3121a 100644
--- a/pkg/provider/provider.go
+++ b/pkg/provider/provider.go
@@ -78,10 +78,16 @@ var providerFlags = []cli.Flag{
 	},
 	&cli.StringFlag{
 		Name:    "update-interval",
-		Aliases: []string{"ui"},
-		Usage:   "Time to wait between distance update checks. The value is an integer string ending in s, m, h for seconds. minutes, hours. Updates will only be seen as fast as they become visible at the upstream location.",
+		Aliases: []string{"uin"},
+		Usage:   "Time to wait between distance update checks when using --follow-dist. The value is an integer string ending in s, m, h for seconds. minutes, hours. Updates will only be seen as fast as they become visible at the upstream location.",
 		Value:   "2m",
 	},
+	&cli.StringFlag{
+		Name:    "update-timeout",
+		Aliases: []string{"uto"},
+		Usage:   "Timeout for getting a provider distance, when using --follow-dist. The value is an integer string ending in s, m, h for seconds. minutes, hours.",
+		Value:   "5m",
+	},
 	&cli.Int64Flag{
 		Name:    "ad-depth-limit",
 		Aliases: []string{"adl"},
@@ -266,9 +272,18 @@ func followDistance(cctx *cli.Context, include, exclude map[peer.ID]struct{}, pc
 		return err
 	}
 
+	var timeout time.Duration
+	updateTimeout := cctx.String("update-timeout")
+	if updateTimeout != "" {
+		timeout, err = time.ParseDuration(updateTimeout)
+		if err != nil {
+			return err
+		}
+	}
+
 	fmt.Fprintln(os.Stderr, "Showing provider distance updates, ctrl-c to cancel...")
 	limit := cctx.Int64("ad-depth-limit")
-	updates := dtrack.RunDistanceTracker(cctx.Context, include, exclude, pc, limit, trackUpdateIn)
+	updates := dtrack.RunDistanceTracker(cctx.Context, include, exclude, pc, limit, trackUpdateIn, timeout)
 	for update := range updates {
 		if update.Err != nil {
 			fmt.Fprintln(os.Stderr, "Provider", update.ID, "distance error:", update.Err)
diff --git a/version.json b/version.json
index c8da44d..59c5faf 100644
--- a/version.json
+++ b/version.json
@@ -1,3 +1,3 @@
 {
-  "version": "v0.0.7"
+  "version": "v0.0.8"
 }