Skip to content

Commit

Permalink
Add traversal.Operation.AddNode
Browse files Browse the repository at this point in the history
  • Loading branch information
anacrolix committed Nov 12, 2023
1 parent 47d54e4 commit 56b782b
Showing 1 changed file with 24 additions and 8 deletions.
32 changes: 24 additions & 8 deletions traversal/operation.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package traversal

import (
"context"
"errors"
"sync/atomic"

"github.com/anacrolix/chansync"
Expand Down Expand Up @@ -121,20 +122,35 @@ func (op *Operation) Stalled() events.Active {
return op.stalled.Active()
}

func (op *Operation) addNodeLocked(n types.AddrMaybeId) (err error) {
if _, ok := op.queried[addrString(n.Addr.String())]; ok {
err = errors.New("already queried")
return
}
if !op.input.NodeFilter(n) {
err = errors.New("failed filter")
return
}
op.unqueried = op.unqueried.Add(n)
op.cond.Broadcast()
return nil
}

// Add an unqueried node returning an error with a reason why the node wasn't added.
func (op *Operation) AddNode(n types.AddrMaybeId) (err error) {
op.mu.Lock()
defer op.mu.Unlock()
return op.addNodeLocked(n)
}

// Add a bunch of unqueried nodes at once, returning how many were successfully added.
func (op *Operation) AddNodes(nodes []types.AddrMaybeId) (added int) {
op.mu.Lock()
defer op.mu.Unlock()
before := op.unqueried.Len()
for _, n := range nodes {
if _, ok := op.queried[addrString(n.Addr.String())]; ok {
continue
}
if !op.input.NodeFilter(n) {
continue
}
op.unqueried = op.unqueried.Add(n)
_ = op.addNodeLocked(n)
}
op.cond.Broadcast()
return op.unqueried.Len() - before
}

Expand Down

0 comments on commit 56b782b

Please sign in to comment.