Add event stream module

antireplay/noop.go

@@ -5,6 +5,7 @@ import "github.com/9seconds/mtg/v2/mtglib"
 type noop struct{}
 
 func (n noop) SeenBefore(_ []byte) bool { return false }
+func (n noop) Shutdown()                {}
 
 func NewNoop() mtglib.AntiReplayCache {
 	return noop{}

antireplay/noop_test.go

@@ -18,6 +18,8 @@ func (suite *NoopTestSuite) TestOp() {
 	suite.False(filter.SeenBefore([]byte{4, 5, 6}))
 	suite.False(filter.SeenBefore([]byte{1, 2, 3}))
 	suite.False(filter.SeenBefore([]byte{4, 5, 6}))
+
+	filter.Shutdown()
 }
 
 func TestNoop(t *testing.T) {

antireplay/stable_bloom_filter.go

@@ -20,6 +20,8 @@ func (s *stableBloomFilter) SeenBefore(digest []byte) bool {
 	return s.filter.TestAndAdd(digest)
 }
 
+func (s *stableBloomFilter) Shutdown() {}
+
 func NewStableBloomFilter(byteSize uint, errorRate float64) mtglib.AntiReplayCache {
 	sf := boom.NewDefaultStableBloomFilter(byteSize*8, errorRate) // nolint: gomnd
 	sf.SetHash(xxhash.New64())

antireplay/stable_bloom_filter_test.go

@@ -18,6 +18,8 @@ func (suite *StableBloomFilterTestSuite) TestOp() {
 	suite.False(filter.SeenBefore([]byte{4, 5, 6}))
 	suite.True(filter.SeenBefore([]byte{1, 2, 3}))
 	suite.True(filter.SeenBefore([]byte{4, 5, 6}))
+
+	filter.Shutdown()
 }
 
 func TestStableBloomFilter(t *testing.T) {

events/event_stream.go

@@ -0,0 +1,64 @@
+package events
+
+import (
+	"context"
+	"runtime"
+
+	"github.com/9seconds/mtg/v2/mtglib"
+	"github.com/OneOfOne/xxhash"
+)
+
+type eventStream struct {
+	ctx       context.Context
+	ctxCancel context.CancelFunc
+	chans     []chan mtglib.Event
+}
+
+func (e eventStream) Send(ctx context.Context, evt mtglib.Event) {
+	chanNo := int(xxhash.ChecksumString32(evt.ConnectionID())) % len(e.chans)
+
+	select {
+	case <-ctx.Done():
+	case <-e.ctx.Done():
+	case e.chans[chanNo] <- evt:
+	}
+}
+
+func (e eventStream) Shutdown() {
+	e.ctxCancel()
+}
+
+func NewEventStream(observerFactories []ObserverFactory) mtglib.EventStream {
+	ctx, cancel := context.WithCancel(context.Background())
+	rv := eventStream{
+		ctx:       ctx,
+		ctxCancel: cancel,
+		chans:     make([]chan mtglib.Event, runtime.NumCPU()),
+	}
+
+	for i := 0; i < runtime.NumCPU(); i++ {
+		rv.chans[i] = make(chan mtglib.Event, 1)
+
+		go eventStreamProcessor(ctx, rv.chans[i], newMultiObserver(observerFactories))
+	}
+
+	return rv
+}
+
+func eventStreamProcessor(ctx context.Context, eventChan <-chan mtglib.Event, observer Observer) {
+	defer observer.Shutdown()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case evt := <-eventChan:
+			switch typedEvt := evt.(type) {
+			case mtglib.EventStart:
+				observer.EventStart(typedEvt)
+			case mtglib.EventFinish:
+				observer.EventFinish(typedEvt)
+			}
+		}
+	}
+}

events/init.go

@@ -0,0 +1,12 @@
+package events
+
+import "github.com/9seconds/mtg/v2/mtglib"
+
+type Observer interface {
+	EventStart(mtglib.EventStart)
+	EventFinish(mtglib.EventFinish)
+
+	Shutdown()
+}
+
+type ObserverFactory func() Observer

events/multi_observer.go

@@ -0,0 +1,59 @@
+package events
+
+import (
+	"sync"
+
+	"github.com/9seconds/mtg/v2/mtglib"
+)
+
+type multiObserver struct {
+	observers []Observer
+}
+
+func (m multiObserver) EventStart(evt mtglib.EventStart) {
+	wg := &sync.WaitGroup{}
+	wg.Add(len(m.observers))
+
+	for _, v := range m.observers {
+		go func(obs Observer) {
+			defer wg.Done()
+
+			obs.EventStart(evt)
+		}(v)
+	}
+
+	wg.Wait()
+}
+
+func (m multiObserver) EventFinish(evt mtglib.EventFinish) {
+	wg := &sync.WaitGroup{}
+	wg.Add(len(m.observers))
+
+	for _, v := range m.observers {
+		go func(obs Observer) {
+			defer wg.Done()
+
+			obs.EventFinish(evt)
+		}(v)
+	}
+
+	wg.Wait()
+}
+
+func (m multiObserver) Shutdown() {
+	for _, v := range m.observers {
+		v.Shutdown()
+	}
+}
+
+func newMultiObserver(factories []ObserverFactory) Observer {
+	observers := make([]Observer, len(factories))
+
+	for i, v := range factories {
+		observers[i] = v()
+	}
+
+	return multiObserver{
+		observers: observers,
+	}
+}

ipblocklist/firehol.go

@@ -29,14 +29,20 @@ const (
 var fireholRegexpComment = regexp.MustCompile(`\s*#.*?$`)
 
 type Firehol struct {
-	logger     mtglib.Logger
+	ctx       context.Context
-	rwMutex    sync.RWMutex
+	ctxCancel context.CancelFunc
+	logger    mtglib.Logger
+
+	rwMutex sync.RWMutex
+
 	remoteURLs []string
 	localFiles []string
+
 	httpClient *http.Client
 	workerPool *ants.Pool
-	treeV4     *bool_tree.TreeV4
+
-	treeV6     *bool_tree.TreeV6
+	treeV4 *bool_tree.TreeV4
+	treeV6 *bool_tree.TreeV6
 }
 
 func (f *Firehol) Contains(ip net.IP) bool {
@@ -76,7 +82,7 @@ func (f *Firehol) containsIPv6(addr net.IP) bool {
 	return false
 }
 
-func (f *Firehol) Run(ctx context.Context, updateEach time.Duration) {
+func (f *Firehol) Run(updateEach time.Duration) {
 	ticker := time.NewTicker(updateEach)
 
 	defer func() {
@@ -88,24 +94,28 @@ func (f *Firehol) Run(ctx context.Context, updateEach time.Duration) {
 		}
 	}()
 
-	if err := f.update(ctx); err != nil {
+	if err := f.update(); err != nil {
 		f.logger.WarningError("cannot update blocklist", err)
 	}
 
 	for {
 		select {
-		case <-ctx.Done():
+		case <-f.ctx.Done():
 			return
 		case <-ticker.C:
-			if err := f.update(ctx); err != nil {
+			if err := f.update(); err != nil {
 				f.logger.WarningError("cannot update blocklist", err)
 			}
 		}
 	}
 }
 
-func (f *Firehol) update(ctx context.Context) error { // nolint: funlen, cyclop
+func (f *Firehol) Shutdown() {
-	ctx, cancel := context.WithCancel(ctx)
+	f.ctxCancel()
+}
+
+func (f *Firehol) update() error { // nolint: funlen, cyclop
+	ctx, cancel := context.WithCancel(f.ctx)
 	defer cancel()
 
 	wg := &sync.WaitGroup{}
@@ -314,8 +324,11 @@ func NewFirehol(logger mtglib.Logger, network mtglib.Network,
 	}
 
 	workerPool, _ := ants.NewPool(int(downloadConcurrency))
+	ctx, cancel := context.WithCancel(context.Background())
 
 	return &Firehol{
+		ctx:        ctx,
+		ctxCancel:  cancel,
 		logger:     logger.Named("firehol"),
 		httpClient: network.MakeHTTPClient(nil),
 		treeV4:     bool_tree.NewTreeV4(),

ipblocklist/firehol_test.go

@@ -1,7 +1,6 @@
 package ipblocklist_test
 
 import (
-	"context"
 	"io"
 	"net"
 	"net/http"
@@ -72,12 +71,15 @@ func (suite *FireholTestSuite) TestLocalFail() {
 
 	suite.NoError(err)
 
-	go blocklist.Run(context.Background(), time.Hour)
+	go blocklist.Run(time.Hour)
 
 	time.Sleep(500 * time.Millisecond)
 
 	suite.False(blocklist.Contains(net.ParseIP("10.0.0.10")))
 	suite.False(blocklist.Contains(net.ParseIP("127.0.0.1")))
+
+	blocklist.Shutdown()
+	time.Sleep(500 * time.Millisecond)
 }
 
 func (suite *FireholTestSuite) TestLocalOk() {
@@ -87,12 +89,15 @@ func (suite *FireholTestSuite) TestLocalOk() {
 
 	suite.NoError(err)
 
-	go blocklist.Run(context.Background(), time.Hour)
+	go blocklist.Run(time.Hour)
 
 	time.Sleep(500 * time.Millisecond)
 
 	suite.True(blocklist.Contains(net.ParseIP("10.0.0.10")))
 	suite.False(blocklist.Contains(net.ParseIP("127.0.0.1")))
+
+	blocklist.Shutdown()
+	time.Sleep(500 * time.Millisecond)
 }
 
 func (suite *FireholTestSuite) TestRemoteFail() {
@@ -102,11 +107,14 @@ func (suite *FireholTestSuite) TestRemoteFail() {
 
 	suite.NoError(err)
 
-	go blocklist.Run(context.Background(), time.Hour)
+	go blocklist.Run(time.Hour)
 
 	time.Sleep(500 * time.Millisecond)
 
 	suite.False(blocklist.Contains(net.ParseIP("10.2.2.2")))
+
+	blocklist.Shutdown()
+	time.Sleep(500 * time.Millisecond)
 }
 
 func (suite *FireholTestSuite) TestMixed() {
@@ -123,12 +131,15 @@ func (suite *FireholTestSuite) TestMixed() {
 
 	suite.NoError(err)
 
-	go blocklist.Run(context.Background(), time.Hour)
+	go blocklist.Run(time.Hour)
 
 	time.Sleep(500 * time.Millisecond)
 
 	suite.True(blocklist.Contains(net.ParseIP("10.2.2.2")))
 	suite.True(blocklist.Contains(net.ParseIP("10.1.0.100")))
+
+	blocklist.Shutdown()
+	time.Sleep(500 * time.Millisecond)
 }
 
 func TestFirehol(t *testing.T) {

ipblocklist/noop.go

@@ -9,6 +9,7 @@ import (
 type noop struct{}
 
 func (n noop) Contains(ip net.IP) bool { return false }
+func (n noop) Shutdown()               {}
 
 func NewNoop() mtglib.IPBlocklist {
 	return noop{}

mtglib/events.go

@@ -0,0 +1,25 @@
+package mtglib
+
+import (
+	"net"
+	"time"
+)
+
+type eventBase struct {
+	CreatedAt time.Time
+	ConnID    string
+}
+
+func (e eventBase) ConnectionID() string {
+	return e.ConnID
+}
+
+type EventStart struct {
+	eventBase
+
+	RemoteIP net.IP
+}
+
+type EventFinish struct {
+	eventBase
+}

mtglib/init.go

@@ -19,10 +19,21 @@ type Network interface {
 
 type AntiReplayCache interface {
 	SeenBefore(data []byte) bool
+	Shutdown()
 }
 
 type IPBlocklist interface {
 	Contains(net.IP) bool
+	Shutdown()
+}
+
+type Event interface {
+	ConnectionID() string
+}
+
+type EventStream interface {
+	Send(context.Context, Event)
+	Shutdown()
 }
 
 type Logger interface {

testlib/events_observer_mock.go

@@ -0,0 +1,7 @@
+package testlib
+
+import "github.com/stretchr/testify/mock"
+
+type EventsObserverMock struct {
+	mock.Mock
+}