Merge pull request #38 from umpc/dev

Simplify/Fix Iter Close method.

iter.go

@@ -15,24 +15,13 @@ type IterChCloser struct {
 // Close cancels a channel-based iteration and causes the sending goroutine to exit.
 // Close should be used after an IterChCloser is finished being read from.
 func (iterCh *IterChCloser) Close() error {
-	select {
-	case iterCh.canceled <- struct{}{}:
-	default:
-	}
-
+	close(iterCh.canceled)
 	return nil
 }
 
-// Records returns nil if the IterChCloser has been closed.
-// Otherwise, Record returns a channel that records can be read from.
+// Records returns a channel that records can be read from.
 func (iterCh *IterChCloser) Records() <-chan Record {
-	select {
-	case <-iterCh.canceled:
-		iterCh.canceled <- struct{}{}
-		return nil
-	default:
-		return iterCh.ch
-	}
+	return iterCh.ch
 }
 
 // IterChParams contains configurable settings for CustomIterCh.
@@ -72,19 +61,21 @@ func (sm *SortedMap) recordFromIdx(i int) Record {
 }
 
 func (sm *SortedMap) sendRecord(iterCh IterChCloser, sendTimeout time.Duration, i int) bool {
-	select {
-	case <-iterCh.canceled:
-		iterCh.canceled <- struct{}{}
-		return false
-	default:
-	}
 
 	if sendTimeout <= time.Duration(0) {
-		iterCh.ch <- sm.recordFromIdx(i)
-		return true
+		select {
+		case <-iterCh.canceled:
+			return false
+
+		case iterCh.ch <- sm.recordFromIdx(i):
+			return true
+		}		
 	}
 
 	select {
+	case <-iterCh.canceled:
+		return false
+
 	case iterCh.ch <- sm.recordFromIdx(i):
 		return true
 
@@ -102,7 +93,7 @@ func (sm *SortedMap) iterCh(params IterChParams) (IterChCloser, error) {
 
 	iterCh := IterChCloser{
 		ch:       make(chan Record, setBufSize(params.BufSize)),
-		canceled: make(chan struct{}, 1),
+		canceled: make(chan struct{}),
 	}
 
 	go func(params IterChParams, iterCh IterChCloser) {

iter_test.go

@@ -206,29 +206,33 @@ func TestBounds(t *testing.T) {
 
 	reversed := false
 
-	iterCh, err := sm.BoundedIterCh(reversed, time.Time{}, unixtime)
-	if err != nil {
-		t.Fatal(err)
-	} else {
-		defer iterCh.Close()
-
-		if err := verifyRecords(iterCh.Records(), reversed); err != nil {
+	func() {
+		iterCh, err := sm.BoundedIterCh(reversed, time.Time{}, unixtime)
+		if err != nil {
 			t.Fatal(err)
-		}
-	}
+		} else {
+			defer iterCh.Close()
 
-	iterCh, err = sm.BoundedIterCh(reversed, obsd, github)
-	if err != nil {
-		t.Fatal(err)
-	} else {
-		defer iterCh.Close()
+			if err := verifyRecords(iterCh.Records(), reversed); err != nil {
+				t.Fatal(err)
+			}
+		}
+	}()
 
-		if err := verifyRecords(iterCh.Records(), reversed); err != nil {
+	func() {
+		iterCh, err := sm.BoundedIterCh(reversed, obsd, github)
+		if err != nil {
 			t.Fatal(err)
+		} else {
+			defer iterCh.Close()
+
+			if err := verifyRecords(iterCh.Records(), reversed); err != nil {
+				t.Fatal(err)
+			}
 		}
-	}
+	}()
 
-	_, err = sm.BoundedIterCh(reversed, obsd, obsd)
+	_, err := sm.BoundedIterCh(reversed, obsd, obsd)
 	if err == nil {
 		t.Fatal("equal bounds values were accepted error")
 	}
@@ -316,7 +320,7 @@ func TestCustomIterCh(t *testing.T) {
 	}()
 }
 
-func TestCancelCustomIterCh(t *testing.T) {
+func TestCloseCustomIterCh(t *testing.T) {
 	sm, _, err := newSortedMapFromRandRecords(1000)
 	if err != nil {
 		t.Fatal(err)
@@ -325,36 +329,25 @@ func TestCancelCustomIterCh(t *testing.T) {
 	earlierDate := time.Date(1, 1, 1, 0, 0, 0, 0, time.UTC)
 	laterDate := time.Now()
 
-	func() {
-		params := IterChParams{
-			LowerBound: earlierDate,
-			UpperBound: laterDate,
-		}
+	ch, err := sm.IterCh()
+	if err != nil {
+		t.Fatal(err)
+	}
 
-		ch, err := sm.CustomIterCh(params)
-		if err != nil {
-			t.Fatal(err)
-		}
-		defer ch.Close()
+	ch.Close()
 
-		go func(ch IterChCloser) {
-			i := 0
-			for range ch.Records() {
-				if i == 50 {
-					ch.Close()
-				}
-				i++
-			}
-			if err := verifyRecords(ch.Records(), params.Reversed); err != nil {
-				if err.Error() != "Channel was nil." {
-					t.Fatal(err)
-				}
-			} else {
-				t.Fatal("Channel was not closed.")
-			}
-		}(ch)
-		ch.Close()
-	}()
+	params := IterChParams{
+		SendTimeout: 5 * time.Minute,
+		LowerBound: earlierDate,
+		UpperBound: laterDate,
+	}
+
+	ch, err = sm.CustomIterCh(params)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	ch.Close()
 }
 
 func TestIterFunc(t *testing.T) {