From b8c847878322572dbffe9b3b3cfce9bf13b896df Mon Sep 17 00:00:00 2001
From: Steve Yen <steve.yen@gmail.com>
Date: Thu, 21 Jul 2016 10:22:06 -0700
Subject: [PATCH] optimize collector to check ctx.Done() only occasionally

---
 search/collectors/collector_top_score.go | 26 ++++++++++++++----------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/search/collectors/collector_top_score.go b/search/collectors/collector_top_score.go
index 740427ec..90b0c75e 100644
--- a/search/collectors/collector_top_score.go
+++ b/search/collectors/collector_top_score.go
@@ -57,6 +57,8 @@ func (tksc *TopScoreCollector) Took() time.Duration {
 	return tksc.took
 }
 
+var COLLECT_CHECK_DONE_EVERY = uint64(1024)
+
 func (tksc *TopScoreCollector) Collect(ctx context.Context, searcher search.Searcher) error {
 	startTime := time.Now()
 	var err error
@@ -69,19 +71,21 @@ func (tksc *TopScoreCollector) Collect(ctx context.Context, searcher search.Sear
 		next, err = searcher.Next(&pre)
 	}
 	for err == nil && next != nil {
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		default:
-			tksc.collectSingle(next)
-			if tksc.facetsBuilder != nil {
-				err = tksc.facetsBuilder.Update(next)
-				if err != nil {
-					break
-				}
+		if tksc.total%COLLECT_CHECK_DONE_EVERY == 0 {
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			default:
 			}
-			next, err = searcher.Next(pre.Reset())
 		}
+		tksc.collectSingle(next)
+		if tksc.facetsBuilder != nil {
+			err = tksc.facetsBuilder.Update(next)
+			if err != nil {
+				break
+			}
+		}
+		next, err = searcher.Next(pre.Reset())
 	}
 	// compute search duration
 	tksc.took = time.Since(startTime)