Fix remove peer from workers heap

protocol/ethereum/backend.go

@@ -74,10 +74,7 @@ func NewBackend(minimal *minimal.Minimal, blockchain *blockchain.Blockchain) (*B
 		waitCh:      make([]chan struct{}, 0),
 
 		// -- new fields
-		heap: &workersHeap{
-			index: map[string]*worker{},
-			heap:  workersHeapImpl{},
-		},
+		heap:      newWorkersHeap(),
 		tasks:     map[uint64]*task{},
 		tasksLock: sync.Mutex{},
 		wakeCh:    make(chan struct{}, 10),
@@ -558,6 +555,13 @@ type worker struct {
 	proto       *Ethereum
 }
 
+func newWorkersHeap() *workersHeap {
+	return &workersHeap{
+		index: map[string]*worker{},
+		heap:  workersHeapImpl{},
+	}
+}
+
 type workersHeap struct {
 	index map[string]*worker
 	heap  workersHeapImpl
@@ -568,6 +572,7 @@ func (w *workersHeap) Remove(id string) {
 	if !ok {
 		return
 	}
+	delete(w.index, id)
 	heap.Remove(&w.heap, wrk.index)
 	return
 }

protocol/ethereum/backend_test.go

@@ -260,3 +260,22 @@ func TestBackendNotify(t *testing.T) {
 	fmt.Println("-- forks --")
 	fmt.Println(b.GetForks())
 }
+
+func TestWorkerPool(t *testing.T) {
+	pool := newWorkersHeap()
+
+	// pool is nil the first time
+	assert.Nil(t, pool.Peek())
+
+	// push and peek
+	assert.NoError(t, pool.Push("1", nil))
+	assert.Equal(t, pool.Peek().id, "1")
+
+	assert.NoError(t, pool.Push("2", nil))
+	assert.Len(t, pool.index, 2)
+
+	// Remove '1'. Only one element left
+	pool.Remove("1")
+	assert.Len(t, pool.index, 1)
+	assert.Equal(t, pool.Peek().id, "2")
+}