// Returns the list of nodes to report to, and the report for each.
func (t *Topic) GenerateReports() ([]*big.Int, []int) {
t.lock.RLock()
defer t.lock.RUnlock()
// Copy child list (rem local if subbed) + parent if not topic root
ids := make([]*big.Int, len(t.nodes), len(t.nodes)+1)
copy(ids, t.nodes)
idx := sortext.SearchBigInts(ids, t.owner)
if idx < len(ids) && t.owner.Cmp(ids[idx]) == 0 {
last := len(ids) - 1
ids[idx] = ids[last]
ids = ids[:last]
}
if t.parent != nil {
ids = append(ids, t.parent)
}
// Calculate the capacities that should be reported to each
caps := make([]int, len(ids))
for i, id := range ids {
caps[i] = t.load.Capacity(id)
}
// Return the capacity with the nodes to report to
return ids, caps
}
// Returns the list of nodes that a broadcast message should be sent to. An
// optional ex node can be specified to exclude it from the list.
func (t *Topic) Broadcast(ex *big.Int) []*big.Int {
t.lock.RLock()
defer t.lock.RUnlock()
// Gather all the nodes to broadcast to
nodes := make([]*big.Int, len(t.nodes), len(t.nodes)+1)
copy(nodes, t.nodes)
if t.parent != nil {
nodes = append(nodes, t.parent)
}
// If exclusion is needed, do it
if ex != nil {
// Sort the nodes and do a binary search on them
sortext.BigInts(nodes)
idx := sortext.SearchBigInts(nodes, ex)
// Swap out with the last if found
if idx < len(nodes) && ex.Cmp(nodes[idx]) == 0 {
last := len(nodes) - 1
nodes[idx] = nodes[last]
nodes = nodes[:last]
}
}
return nodes
}
// If local subscriptions are alive in the topic, updates the balancer according
// to the messages processed since the last beat.
func (t *Topic) Cycle() {
t.lock.RLock()
defer t.lock.RUnlock()
// Notify the balancer of the local capacity
idx := sortext.SearchBigInts(t.nodes, t.owner)
if idx < len(t.nodes) && t.owner.Cmp(t.nodes[idx]) == 0 {
// Sanity check not to send some weird value
cap := math.Max(0, float64(atomic.LoadInt32(&t.msgs))/float64(system.CpuUsage()))
cap = math.Min(math.MaxInt32, cap)
t.load.Update(t.owner, int(cap))
}
// Reset counters for next beat
atomic.StoreInt32(&t.msgs, 0)
}
// Subscribes a node to the topic and inserts it into the load balancer registry.
func (t *Topic) Subscribe(id *big.Int) error {
t.lock.Lock()
defer t.lock.Unlock()
// log.Printf("%v:%v: node subscription: %v.", t.owner, t.id, id)
// Ensure double subscription doesn't happen
idx := sortext.SearchBigInts(t.nodes, id)
if idx < len(t.nodes) && id.Cmp(t.nodes[idx]) == 0 {
return ErrSubscribed
}
// New entity, insert into the list
t.nodes = append(t.nodes, id)
sortext.BigInts(t.nodes)
t.members[id.String()] = struct{}{}
// log.Printf("%v:%v: subbed, state: %v.", t.owner, t.id, t.nodes)
// Start load balancing to it too
t.load.Register(id)
return nil
}
// Unregisters a node from the topic, removing it from the balancer's registry.
func (t *Topic) Unsubscribe(id *big.Int) error {
t.lock.Lock()
defer t.lock.Unlock()
// log.Printf("%v:%v: node unsubscription: %v.", t.owner, t.id, id)
// Ensure double unsubscription doesn't happen
idx := sortext.SearchBigInts(t.nodes, id)
if idx == len(t.nodes) || id.Cmp(t.nodes[idx]) != 0 {
return ErrNotSubscribed
}
// Remove the node from the children
last := len(t.nodes) - 1
t.nodes[idx] = t.nodes[last]
t.nodes = t.nodes[:last]
sortext.BigInts(t.nodes)
delete(t.members, id.String())
// log.Printf("%v:%v: remed, state: %v.", t.owner, t.id, t.nodes)
// Remove the node from the load balancer
t.load.Unregister(id)
return nil
}
func TestTopic(t *testing.T) {
// Define some setup parameters for the test
topicId := big.NewInt(314)
ownerId := big.NewInt(141)
nodeIds := []int64{1, 2, 3, 4, 5}
nodes := make([]*big.Int, len(nodeIds))
for i, id := range nodeIds {
nodes[i] = big.NewInt(id)
}
sortext.BigInts(nodes)
// Create the topic and check internal state
top := New(topicId, ownerId)
if id := top.Self(); id.Cmp(topicId) != 0 {
t.Fatalf("topic id mismatch: have %v, want %v.", id, topicId)
}
if id := top.owner; id.Cmp(ownerId) != 0 {
t.Fatalf("topic owner mismatch: have %v, want %v.", id, ownerId)
}
// Check subscribe and unsubscribe
for i, id := range nodes {
top.Subscribe(id)
if n := len(top.nodes); n != i+1 {
t.Fatalf("topic child node count mismatch: have %v, want %v", n, i+1)
}
}
for i, id := range nodes {
top.Unsubscribe(id)
if n := len(top.nodes); n != len(nodes)-1-i {
t.Fatalf("topic child node count mismatch: have %v, want %v", n, len(nodes)-1-i)
}
}
// Subscribe everybody back
for _, id := range nodes {
top.Subscribe(id)
}
// Check broadcasting with no exclusion
ns := top.Broadcast(nil)
if len(ns) != len(nodes) {
t.Fatalf("broadcast node list length mismatch: have %v, want %v.", len(ns), len(nodes))
}
for i, id := range ns {
if nodes[i].Cmp(id) != 0 {
t.Fatalf("broadcast node %d mismatch: have %v, want %v.", i, id, nodes[i])
}
}
// Check broadcasting with exclusion
ns = top.Broadcast(nodes[0])
if len(ns) != len(nodes)-1 {
t.Fatalf("excluded broadcast node list length mismatch: have %v, want %v.", len(ns), len(nodes)-1)
}
for _, id := range ns {
if nodes[0].Cmp(id) == 0 {
t.Fatalf("broadcast includes excluded node: %v.", id)
}
}
// Check load balancing (without one entry)
ns = []*big.Int{}
for i := 0; i < 1000; i++ {
n, err := top.Balance(nodes[0])
if err != nil {
t.Fatalf("failed to balance: %v.", err)
}
ns = append(ns, n)
}
if len(ns) == 0 {
t.Fatalf("no nodes have been balanced to")
}
for i, id := range ns {
if id.Cmp(nodes[0]) == 0 {
t.Fatalf("balance %d: reached excluded node %v.", i, id)
}
idx := sortext.SearchBigInts(nodes, id)
if idx >= len(nodes) || nodes[idx] != id {
t.Fatalf("balance %d: invalid node id %v.", i, id)
}
}
// Add a local subscription
if err := top.Subscribe(ownerId); err != nil {
t.Fatalf("failed to subscribe with local node: %v.", err)
}
// Check load report generation
ns, caps := top.GenerateReports()
if len(ns) != len(nodes) || len(caps) != len(nodes) {
t.Fatalf("report target size mismatch: have %v/%v nodes/caps, want %v.", len(ns), len(caps), len(nodes))
}
for i, cap := range caps {
if cap != len(nodes) {
t.Fatalf("capacity %d mismatch: have %v, want %v", i, cap, len(nodes))
}
}
// Check load processing
total := 1 // Local apps
for i, id := range nodes {
top.ProcessReport(id, 10*(i+1))
total += 10 * (i + 1)
}
ns, caps = top.GenerateReports()
for i, cap := range caps {
if cap != total-10*(i+1) {
t.Fatalf("capacity %d mismatch: have %v, want %v", i, cap, total-10*(i+1))
}
}
}