Remove race conditions when accessing the last chunk

3 files changed, 36 insertions, 20 deletions

diff --git a/src/chunklist.go b/src/chunklist.go
index b1f9638d..5bca6da8 100644
--- a/src/chunklist.go
+++ b/src/chunklist.go
@@ -42,14 +42,6 @@ func CountItems(cs []*Chunk) int {
 	return CHUNK_SIZE*(len(cs)-1) + len(*(cs[len(cs)-1]))
 }
 
-func (cl *ChunkList) Count() int {
-	return cl.count
-}
-
-func (cl *ChunkList) Chunks() []*Chunk {
-	return cl.chunks
-}
-
 func (cl *ChunkList) Push(data string) {
 	cl.mutex.Lock()
 	defer cl.mutex.Unlock()
@@ -63,11 +55,24 @@ func (cl *ChunkList) Push(data string) {
 	cl.count += 1
 }
 
-func (cl *ChunkList) Snapshot() []*Chunk {
+func (cl *ChunkList) Snapshot() ([]*Chunk, int) {
 	cl.mutex.Lock()
 	defer cl.mutex.Unlock()
 
 	ret := make([]*Chunk, len(cl.chunks))
 	copy(ret, cl.chunks)
-	return ret
+
+	// Duplicate the last chunk
+	if cnt := len(ret); cnt > 0 {
+		ret[cnt-1] = ret[cnt-1].dupe()
+	}
+	return ret, cl.count
+}
+
+func (c *Chunk) dupe() *Chunk {
+	newChunk := make(Chunk, len(*c))
+	for idx, ptr := range *c {
+		newChunk[idx] = ptr
+	}
+	return &newChunk
 }
diff --git a/src/chunklist_test.go b/src/chunklist_test.go
index a7daa47e..b244ece8 100644
--- a/src/chunklist_test.go
+++ b/src/chunklist_test.go
@@ -11,8 +11,8 @@ func TestChunkList(t *testing.T) {
 	})
 
 	// Snapshot
-	snapshot := cl.Snapshot()
-	if len(snapshot) > 0 {
+	snapshot, count := cl.Snapshot()
+	if len(snapshot) > 0 || count > 0 {
 		t.Error("Snapshot should be empty now")
 	}
 
@@ -26,8 +26,8 @@ func TestChunkList(t *testing.T) {
 	}
 
 	// But the new snapshot should contain the added items
-	snapshot = cl.Snapshot()
-	if len(snapshot) != 1 {
+	snapshot, count = cl.Snapshot()
+	if len(snapshot) != 1 && count != 2 {
 		t.Error("Snapshot should not be empty now")
 	}
 
@@ -55,12 +55,20 @@ func TestChunkList(t *testing.T) {
 	}
 
 	// New snapshot
-	snapshot = cl.Snapshot()
+	snapshot, count = cl.Snapshot()
 	if len(snapshot) != 3 || !snapshot[0].IsFull() ||
-		!snapshot[1].IsFull() || snapshot[2].IsFull() {
+		!snapshot[1].IsFull() || snapshot[2].IsFull() || count != CHUNK_SIZE*2+2 {
 		t.Error("Expected two full chunks and one more chunk")
 	}
 	if len(*snapshot[2]) != 2 {
 		t.Error("Unexpected number of items")
 	}
+
+	cl.Push("hello")
+	cl.Push("world")
+
+	lastChunkCount := len(*snapshot[len(snapshot)-1])
+	if lastChunkCount != 2 {
+		t.Error("Unexpected number of items:", lastChunkCount)
+	}
 }
diff --git a/src/core.go b/src/core.go
index 7abee80a..b6f08572 100644
--- a/src/core.go
+++ b/src/core.go
@@ -90,8 +90,9 @@ func Run(options *Options) {
 			})
 		}
 
+		snapshot, _ := chunkList.Snapshot()
 		matches, cancelled := matcher.scan(MatchRequest{
-			chunks:  chunkList.Snapshot(),
+			chunks:  snapshot,
 			pattern: pattern}, limit)
 
 		if !cancelled && (filtering ||
@@ -127,11 +128,13 @@ func Run(options *Options) {
 
 				case EVT_READ_NEW, EVT_READ_FIN:
 					reading = reading && evt == EVT_READ_NEW
-					terminal.UpdateCount(chunkList.Count(), !reading)
-					matcher.Reset(chunkList.Snapshot(), terminal.Input(), false)
+					snapshot, count := chunkList.Snapshot()
+					terminal.UpdateCount(count, !reading)
+					matcher.Reset(snapshot, terminal.Input(), false)
 
 				case EVT_SEARCH_NEW:
-					matcher.Reset(chunkList.Snapshot(), terminal.Input(), true)
+					snapshot, _ := chunkList.Snapshot()
+					matcher.Reset(snapshot, terminal.Input(), true)
 					delay = false
 
 				case EVT_SEARCH_PROGRESS: