From 5999bf88ac4ca1d21ad52534b8decb2d6e1f994e Mon Sep 17 00:00:00 2001
From: Michael Avila <davidmichaelavila@gmail.com>
Date: Wed, 12 Jun 2019 13:02:41 -0700
Subject: [PATCH] Either timeout or not

---
 simple/provider.go      | 11 +++++++--
 simple/provider_test.go | 50 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/simple/provider.go b/simple/provider.go
index 9b374a0..9302a27 100644
--- a/simple/provider.go
+++ b/simple/provider.go
@@ -88,8 +88,15 @@ func (p *Provider) handleAnnouncements() {
 				case <-p.ctx.Done():
 					return
 				case c := <-p.queue.Dequeue():
-					ctx, cancel := context.WithTimeout(p.ctx, p.timeout)
-					defer cancel()
+					var ctx context.Context
+					var cancel context.CancelFunc
+					if p.timeout > 0 {
+						ctx, cancel = context.WithTimeout(p.ctx, p.timeout)
+						defer cancel()
+					} else {
+						ctx = p.ctx
+					}
+
 					logP.Info("announce - start - ", c)
 					if err := p.contentRouting.Provide(ctx, c, true); err != nil {
 						logP.Warningf("Unable to provide entry: %s, %s", c, err)
diff --git a/simple/provider_test.go b/simple/provider_test.go
index 6fbc528..71a299a 100644
--- a/simple/provider_test.go
+++ b/simple/provider_test.go
@@ -24,7 +24,11 @@ type mockRouting struct {
 }
 
 func (r *mockRouting) Provide(ctx context.Context, cid cid.Cid, recursive bool) error {
-	r.provided <- cid
+	select {
+	case r.provided <- cid:
+	case <-ctx.Done():
+		panic("context cancelled, but shouldn't have")
+	}
 	return nil
 }
 
@@ -81,3 +85,47 @@ func TestAnnouncement(t *testing.T) {
 		}
 	}
 }
+
+func TestAnnouncementTimeout(t *testing.T) {
+	ctx := context.Background()
+	defer ctx.Done()
+
+	ds := sync.MutexWrap(datastore.NewMapDatastore())
+	queue, err := q.NewQueue(ctx, "test", ds)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	r := mockContentRouting()
+
+	prov := NewProvider(ctx, queue, r, WithTimeout(1 * time.Second))
+	prov.Run()
+
+	cids := cid.NewSet()
+
+	for i := 0; i < 100; i++ {
+		c := blockGenerator.Next().Cid()
+		cids.Add(c)
+	}
+
+	go func() {
+		for _, c := range cids.Keys() {
+			err = prov.Provide(c)
+			// A little goroutine stirring to exercise some different states
+			r := rand.Intn(10)
+			time.Sleep(time.Microsecond * time.Duration(r))
+		}
+	}()
+
+	for cids.Len() > 0 {
+		select {
+		case cp := <-r.provided:
+			if !cids.Has(cp) {
+				t.Fatal("Wrong CID provided")
+			}
+			cids.Remove(cp)
+		case <-time.After(time.Second * 5):
+			t.Fatal("Timeout waiting for cids to be provided.")
+		}
+	}
+}