kubo/thirdparty/notifier/notifier_test.go

package notifier

import (
	"fmt"
	"sync"
	"testing"
)

// test data structures
type Router struct {
	queue    chan Packet
	notifier Notifier
}

type Packet struct{}

type RouterNotifiee interface {
	Enqueued(*Router, Packet)
	Forwarded(*Router, Packet)
	Dropped(*Router, Packet)
}

func (r *Router) Notify(n RouterNotifiee) {
	r.notifier.Notify(n)
}

func (r *Router) StopNotify(n RouterNotifiee) {
	r.notifier.StopNotify(n)
}

func (r *Router) notifyAll(notify func(n RouterNotifiee)) {
	r.notifier.NotifyAll(func(n Notifiee) {
		notify(n.(RouterNotifiee))
	})
}

func (r *Router) Receive(p Packet) {

	select {
	case r.queue <- p: // enqueued
		r.notifyAll(func(n RouterNotifiee) {
			n.Enqueued(r, p)
		})

	default: // drop
		r.notifyAll(func(n RouterNotifiee) {
			n.Dropped(r, p)
		})
	}
}

func (r *Router) Forward() {
	p := <-r.queue
	r.notifyAll(func(n RouterNotifiee) {
		n.Forwarded(r, p)
	})
}

type Metrics struct {
	enqueued  int
	forwarded int
	dropped   int
	received  chan struct{}
	sync.Mutex
}

func (m *Metrics) Enqueued(*Router, Packet) {
	m.Lock()
	m.enqueued++
	m.Unlock()
	if m.received != nil {
		m.received <- struct{}{}
	}
}

func (m *Metrics) Forwarded(*Router, Packet) {
	m.Lock()
	m.forwarded++
	m.Unlock()
	if m.received != nil {
		m.received <- struct{}{}
	}
}

func (m *Metrics) Dropped(*Router, Packet) {
	m.Lock()
	m.dropped++
	m.Unlock()
	if m.received != nil {
		m.received <- struct{}{}
	}
}

func (m *Metrics) String() string {
	m.Lock()
	defer m.Unlock()
	return fmt.Sprintf("%d enqueued, %d forwarded, %d in queue, %d dropped",
		m.enqueued, m.forwarded, m.enqueued-m.forwarded, m.dropped)
}

func TestNotifies(t *testing.T) {

	m := Metrics{received: make(chan struct{})}
	r := Router{queue: make(chan Packet, 10)}
	r.Notify(&m)

	for i := 0; i < 10; i++ {
		r.Receive(Packet{})
		<-m.received
		if m.enqueued != (1 + i) {
			t.Error("not notifying correctly", m.enqueued, 1+i)
		}

	}

	for i := 0; i < 10; i++ {
		r.Receive(Packet{})
		<-m.received
		if m.enqueued != 10 {
			t.Error("not notifying correctly", m.enqueued, 10)
		}
		if m.dropped != (1 + i) {
			t.Error("not notifying correctly", m.dropped, 1+i)
		}
	}
}

func TestStopsNotifying(t *testing.T) {
	m := Metrics{received: make(chan struct{})}
	r := Router{queue: make(chan Packet, 10)}
	r.Notify(&m)

	for i := 0; i < 5; i++ {
		r.Receive(Packet{})
		<-m.received
		if m.enqueued != (1 + i) {
			t.Error("not notifying correctly")
		}
	}

	r.StopNotify(&m)

	for i := 0; i < 5; i++ {
		r.Receive(Packet{})
		select {
		case <-m.received:
			t.Error("did not stop notifying")
		default:
		}
		if m.enqueued != 5 {
			t.Error("did not stop notifying")
		}
	}
}

func TestThreadsafe(t *testing.T) {
	N := 1000
	r := Router{queue: make(chan Packet, 10)}
	m1 := Metrics{received: make(chan struct{})}
	m2 := Metrics{received: make(chan struct{})}
	m3 := Metrics{received: make(chan struct{})}
	r.Notify(&m1)
	r.Notify(&m2)
	r.Notify(&m3)

	var n int
	var wg sync.WaitGroup
	for i := 0; i < N; i++ {
		n++
		wg.Add(1)
		go func() {
			defer wg.Done()
			r.Receive(Packet{})
		}()

		if i%3 == 0 {
			n++
			wg.Add(1)
			go func() {
				defer wg.Done()
				r.Forward()
			}()
		}
	}

	// drain queues
	for i := 0; i < (n * 3); i++ {
		select {
		case <-m1.received:
		case <-m2.received:
		case <-m3.received:
		}
	}

	wg.Wait()

	// counts should be correct and all agree. and this should
	// run fine under `go test -race -cpu=5`

	t.Log("m1", m1.String())
	t.Log("m2", m2.String())
	t.Log("m3", m3.String())

	if m1.String() != m2.String() || m2.String() != m3.String() {
		t.Error("counts disagree")
	}
}