Simplify canceled channel

iter.go

@@ -3,38 +3,24 @@ package sortedmap
 import (
 	"errors"
 	"time"
-	"sync"
 )
 
 // IterChCloser allows records to be read through a channel that is returned by the Records method.
 // IterChCloser values should be closed after use using the Close method.
 type IterChCloser struct {
-	sync.RWMutex
 	ch       chan Record
-	canceled bool
+	canceled chan 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 {
-	iterCh.Lock()
-	iterCh.canceled = true
-	iterCh.Unlock()
-
+	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 {
-	iterCh.RLock()
-	canceled := iterCh.canceled
-	iterCh.RUnlock()
-
-	if canceled {
-		return nil
-	}
-
 	return iterCh.ch
 }
 
@@ -75,20 +61,21 @@ func (sm *SortedMap) recordFromIdx(i int) Record {
 }
 
 func (sm *SortedMap) sendRecord(iterCh IterChCloser, sendTimeout time.Duration, i int) bool {
-	iterCh.RLock()
-	canceled := iterCh.canceled
-	iterCh.RUnlock()
-
-	if canceled {
-		return false
-	}
 
 	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
 
@@ -105,7 +92,8 @@ func (sm *SortedMap) iterCh(params IterChParams) (IterChCloser, error) {
 	}
 
 	iterCh := IterChCloser{
-		ch: make(chan Record, setBufSize(params.BufSize)),
+		ch:       make(chan Record, setBufSize(params.BufSize)),
+		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)
@@ -335,24 +339,8 @@ func TestCancelCustomIterCh(t *testing.T) {
 		if err != nil {
 			t.Fatal(err)
 		}
-		defer 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()
 	}()
 }