Refactor confusing & fragile event bus logic to use callbacks and not channels; ref pterodactyl/panel#2298

events/events.go

@@ -2,6 +2,8 @@ package events
 
 import (
 	"encoding/json"
+	"github.com/pkg/errors"
+	"reflect"
 	"strings"
 	"sync"
 )
@@ -12,14 +14,13 @@ type Event struct {
 }
 
 type EventBus struct {
-	sync.RWMutex
-
-	subscribers map[string]map[chan Event]struct{}
+	mu        sync.RWMutex
+	callbacks map[string][]*func(Event)
 }
 
 func New() *EventBus {
 	return &EventBus{
-		subscribers: make(map[string]map[chan Event]struct{}),
+		callbacks: make(map[string][]*func(Event)),
 	}
 }
 
@@ -39,29 +40,27 @@ func (e *EventBus) Publish(topic string, data string) {
 		}
 	}
 
+	e.mu.RLock()
+	defer e.mu.RUnlock()
+
 	// Acquire a read lock and loop over all of the channels registered for the topic. This
 	// avoids a panic crash if the process tries to unregister the channel while this routine
 	// is running.
-	go func() {
-		e.RLock()
-		defer e.RUnlock()
-
-		if ch, ok := e.subscribers[t]; ok {
-			e := Event{Data: data, Topic: topic}
-
-			for channel := range ch {
-				go func(channel chan Event, e Event) {
-					channel <- e
-				}(channel, e)
-			}
+	if _, ok := e.callbacks[t]; ok {
+		evt := Event{Data: data, Topic: topic}
+		for _, callback := range e.callbacks[t] {
+			go func(evt Event, callback func(Event)) {
+				callback(evt)
+			}(evt, *callback)
 		}
-	}()
+	}
 }
 
+// Publishes a JSON message to a given topic.
 func (e *EventBus) PublishJson(topic string, data interface{}) error {
 	b, err := json.Marshal(data)
 	if err != nil {
-		return err
+		return errors.WithStack(err)
 	}
 
 	e.Publish(topic, string(b))
@@ -69,45 +68,65 @@ func (e *EventBus) PublishJson(topic string, data interface{}) error {
 	return nil
 }
 
-// Subscribe to an emitter topic using a channel.
-func (e *EventBus) Subscribe(topics []string, ch chan Event) {
-	e.Lock()
-	defer e.Unlock()
+// Register a callback function that will be executed each time one of the events using the topic
+// name is called.
+func (e *EventBus) On(topic string, callback *func(Event)) {
+	e.mu.Lock()
+	defer e.mu.Unlock()
 
-	for _, topic := range topics {
-		if _, exists := e.subscribers[topic]; !exists {
-			e.subscribers[topic] = make(map[chan Event]struct{})
-		}
+	// Check if this topic has been registered at least once for the event listener, and if
+	// not create an empty struct for the topic.
+	if _, exists := e.callbacks[topic]; !exists {
+		e.callbacks[topic] = make([]*func(Event), 0)
+	}
 
-		// Only set the channel if there is not currently a matching one for this topic. This
-		// avoids registering two identical listeners for the same topic and causing pain in
-		// the unsubscribe functionality as well.
-		if _, exists := e.subscribers[topic][ch]; !exists {
-			e.subscribers[topic][ch] = struct{}{}
-		}
+	// If this callback is not already registered as an event listener, go ahead and append
+	// it to the array of callbacks for this topic.
+	if e.index(topic, reflect.ValueOf(callback)) < 0 {
+		e.callbacks[topic] = append(e.callbacks[topic], callback)
 	}
 }
 
-// Unsubscribe a channel from a given topic.
-func (e *EventBus) Unsubscribe(topics []string, ch chan Event) {
-	e.Lock()
-	defer e.Unlock()
+// Removes an event listener from the bus.
+func (e *EventBus) Off(topic string, callback *func(Event)) {
+	e.mu.Lock()
+	defer e.mu.Unlock()
 
-	for _, topic := range topics {
-		if _, exists := e.subscribers[topic][ch]; exists {
-			delete(e.subscribers[topic], ch)
+	i := e.index(topic, reflect.ValueOf(callback))
+
+	// If i < 0 it means there was no index found for the given callback, meaning it was
+	// never registered or was already unregistered from the listeners. Also double check
+	// that we didn't somehow escape the length of the topic callback (not sure how that
+	// would happen, but lets avoid a panic condition).
+	if i < 0 || i >= len(e.callbacks[topic]) {
+		return
+	}
+
+	// We can assume that the topic still exists at this point since we acquire an exclusive
+	// lock on the process, and the "e.index" function cannot return a value >= 0 if there is
+	// no topic already existing.
+	e.callbacks[topic] = append(e.callbacks[topic][:i], e.callbacks[topic][i+1:]...)
+}
+
+// Removes all of the event listeners that have been registered for any topic.
+func (e *EventBus) RemoveAll() {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+
+	e.callbacks = make(map[string][]*func(Event))
+}
+
+// Finds the index of a given callback in the topic by comparing all of the registered callback
+// pointers to the passed function. This function does not aquire a lock as it should only be called
+// within the confines of a function that has already acquired a lock for the duration of the lookup.
+func (e *EventBus) index(topic string, v reflect.Value) int {
+	if _, ok := e.callbacks[topic]; ok {
+		for i, handler := range e.callbacks[topic] {
+			if reflect.ValueOf(handler).Pointer() == v.Pointer() {
+				return i
+			}
 		}
 	}
-}
 
-// Removes all of the event listeners for the server. This is used when a server
-// is being deleted to avoid a bunch of de-reference errors cropping up. Obviously
-// should also check elsewhere and handle a server reference going nil, but this
-// won't hurt.
-func (e *EventBus) UnsubscribeAll() {
-	e.Lock()
-	defer e.Unlock()
-
-	// Reset the entire struct into an empty map.
-	e.subscribers = make(map[string]map[chan Event]struct{})
+	return -1
 }