initial commit with basic pub/sub

.gitignore

···
-# If you prefer the allow list template instead of the deny list, see community template:
-# https://github.com/github/gitignore/blob/main/community/Golang/Go.AllowList.gitignore
-#
-# Binaries for programs and plugins
-*.exe
-*.exe~
-*.dll
-*.so
-*.dylib
-
-# Test binary, built with `go test -c`
-*.test
-
-# Output of the go coverage tool, specifically when used with LiteIDE
-*.out
-
-# Dependency directories (remove the comment below to include it)
-# vendor/
-
-# Go workspace file
-go.work
+.DS_STORE

go.mod

···
+module github.com/willdot/message-broker
+
+go 1.21.0
+
+require github.com/stretchr/testify v1.8.4
+
+require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+)

go.sum

···
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
+github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

server.go

···
+package main
+
+import (
+	"context"
+	"encoding/binary"
+	"encoding/json"
+	"fmt"
+	"log/slog"
+	"net"
+	"strings"
+	"sync"
+)
+
+// Action represents the type of action that a connection requests to do
+type Action uint8
+
+const (
+	Subscribe   Action = 1
+	Unsubscribe Action = 2
+	Publish     Action = 3
+)
+
+type message struct {
+	Topic string `json:"topic"`
+	Data  []byte `json:"data"`
+}
+
+type Server struct {
+	addr string
+	lis  net.Listener
+
+	mu     sync.Mutex
+	topics map[string]topic
+}
+
+func NewServer(ctx context.Context, addr string) (*Server, error) {
+	lis, err := net.Listen("tcp", addr)
+	if err != nil {
+		return nil, fmt.Errorf("failed to listen: %w", err)
+	}
+
+	srv := &Server{
+		lis:    lis,
+		topics: map[string]topic{},
+	}
+
+	go srv.start(ctx)
+
+	return srv, nil
+}
+
+func (s *Server) Shutdown() error {
+	return s.lis.Close()
+}
+
+func (s *Server) start(ctx context.Context) {
+	for {
+		conn, err := s.lis.Accept()
+		if err != nil {
+			slog.Error("listener failed to accept", "error", err)
+			// TODO: see if there's a better way to check for this error
+			if strings.Contains(err.Error(), "use of closed network connection") {
+				return
+			}
+		}
+
+		go s.handleConn(conn)
+	}
+}
+
+func getActionFromConn(conn net.Conn) (Action, error) {
+	var action Action
+	err := binary.Read(conn, binary.BigEndian, &action)
+	if err != nil {
+		return 0, err
+	}
+
+	return action, nil
+}
+
+func getDataLengthFromConn(conn net.Conn) (uint32, error) {
+	var dataLen uint32
+	err := binary.Read(conn, binary.BigEndian, &dataLen)
+	if err != nil {
+		return 0, fmt.Errorf("failed to read data length from conn: %w", err)
+	}
+
+	return dataLen, nil
+}
+
+func (s *Server) handleConn(conn net.Conn) {
+	action, err := getActionFromConn(conn)
+	if err != nil {
+		slog.Error("failed to read action from conn", "error", err, "conn", conn.LocalAddr())
+		return
+	}
+
+	switch action {
+	case Subscribe:
+		s.handleSubscribingConn(conn)
+	case Unsubscribe:
+		s.handleUnsubscribingConn(conn)
+	case Publish:
+		s.handlePublisherConn(conn)
+	default:
+		slog.Error("unknown action", "action", action, "conn", conn.LocalAddr())
+		_, _ = conn.Write([]byte("unknown action"))
+	}
+}
+
+func (s *Server) handleSubscribingConn(conn net.Conn) {
+	// subscribe the connection to the topic
+	s.subscribeConnToTopic(conn)
+
+	// keep handling the connection, getting the action from the conection when it wishes to do something else.
+	// once the connection ends, it will be unsubscribed from all topics and returned
+	for {
+		action, err := getActionFromConn(conn)
+		if err != nil {
+			// TODO: see if there's a way to check if the connection has been ended etc
+			slog.Error("failed to read action from subscriber", "error", err, "conn", conn.LocalAddr())
+
+			s.unsubscribeConnectionFromAllTopics(conn.LocalAddr())
+
+			return
+		}
+
+		switch action {
+		case Subscribe:
+			s.subscribeConnToTopic(conn)
+		case Unsubscribe:
+			s.handleUnsubscribingConn(conn)
+		default:
+			slog.Error("unknown action for subscriber", "action", action, "conn", conn.LocalAddr())
+			continue
+		}
+	}
+}
+
+func (s *Server) subscribeConnToTopic(conn net.Conn) {
+	// get the topics the connection wishes to subscribe to
+	dataLen, err := getDataLengthFromConn(conn)
+	if err != nil {
+		slog.Error(err.Error(), "conn", conn.LocalAddr())
+		_, _ = conn.Write([]byte("invalid data length of topics provided"))
+		return
+	}
+	if dataLen == 0 {
+		_, _ = conn.Write([]byte("data length of topics is 0"))
+		return
+	}
+
+	buf := make([]byte, dataLen)
+	_, err = conn.Read(buf)
+	if err != nil {
+		slog.Error("failed to read subscibers topic data", "error", err, "conn", conn.LocalAddr())
+		_, _ = conn.Write([]byte("failed to read topic data"))
+		return
+	}
+
+	var topics []string
+	err = json.Unmarshal(buf, &topics)
+	if err != nil {
+		slog.Error("failed to unmarshal subscibers topic data", "error", err, "conn", conn.LocalAddr())
+		_, _ = conn.Write([]byte("invalid topic data provided"))
+		return
+	}
+
+	s.subscribeToTopics(conn, topics)
+	_, _ = conn.Write([]byte("subscribed"))
+}
+
+func (s *Server) handleUnsubscribingConn(conn net.Conn) {
+	// get the topics the connection wishes to unsubscribe from
+	dataLen, err := getDataLengthFromConn(conn)
+	if err != nil {
+		slog.Error(err.Error(), "conn", conn.LocalAddr())
+		_, _ = conn.Write([]byte("invalid data length of topics provided"))
+		return
+	}
+	if dataLen == 0 {
+		_, _ = conn.Write([]byte("data length of topics is 0"))
+		return
+	}
+
+	buf := make([]byte, dataLen)
+	_, err = conn.Read(buf)
+	if err != nil {
+		slog.Error("failed to read subscibers topic data", "error", err, "conn", conn.LocalAddr())
+		_, _ = conn.Write([]byte("failed to read topic data"))
+		return
+	}
+
+	var topics []string
+	err = json.Unmarshal(buf, &topics)
+	if err != nil {
+		slog.Error("failed to unmarshal subscibers topic data", "error", err, "conn", conn.LocalAddr())
+		_, _ = conn.Write([]byte("invalid topic data provided"))
+		return
+	}
+
+	s.unsubscribeToTopics(conn, topics)
+
+	_, _ = conn.Write([]byte("unsubscribed"))
+}
+
+func (s *Server) handlePublisherConn(conn net.Conn) {
+	dataLen, err := getDataLengthFromConn(conn)
+	if err != nil {
+		slog.Error(err.Error(), "conn", conn.LocalAddr())
+		_, _ = conn.Write([]byte("invalid data length of data provided"))
+		return
+	}
+	if dataLen == 0 {
+		return
+	}
+
+	buf := make([]byte, dataLen)
+	_, err = conn.Read(buf)
+	if err != nil {
+		_, _ = conn.Write([]byte("failed to read data"))
+		slog.Error("failed to read data from conn", "error", err, "conn", conn.LocalAddr())
+		return
+	}
+
+	var msg message
+	err = json.Unmarshal(buf, &msg)
+	if err != nil {
+		_, _ = conn.Write([]byte("invalid message"))
+		slog.Error("failed to unmarshal data to message", "error", err, "conn", conn.LocalAddr())
+		return
+	}
+
+	topic := s.getTopic(msg.Topic)
+	if topic != nil {
+		topic.sendMessageToSubscribers(msg)
+	}
+}
+
+func (s *Server) subscribeToTopics(conn net.Conn, topics []string) {
+	for _, topic := range topics {
+		s.addSubsciberToTopic(topic, conn)
+	}
+}
+
+func (s *Server) addSubsciberToTopic(topicName string, conn net.Conn) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	t, ok := s.topics[topicName]
+	if !ok {
+		t = newTopic(topicName)
+	}
+
+	t.subscriptions[conn.LocalAddr()] = Subscriber{
+		conn:          conn,
+		currentOffset: 0,
+	}
+
+	s.topics[topicName] = t
+}
+
+func (s *Server) unsubscribeToTopics(conn net.Conn, topics []string) {
+	for _, topic := range topics {
+		s.removeSubsciberFromTopic(topic, conn)
+	}
+}
+
+func (s *Server) removeSubsciberFromTopic(topicName string, conn net.Conn) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	t, ok := s.topics[topicName]
+	if !ok {
+		return
+	}
+
+	delete(t.subscriptions, conn.LocalAddr())
+}
+
+func (s *Server) unsubscribeConnectionFromAllTopics(addr net.Addr) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	for _, topic := range s.topics {
+		delete(topic.subscriptions, addr)
+	}
+}
+
+func (s *Server) getTopic(topicName string) *topic {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if topic, ok := s.topics[topicName]; ok {
+		return &topic
+	}
+
+	return nil
+}

server_test.go

···
+package main
+
+import (
+	"context"
+	"encoding/binary"
+	"encoding/json"
+	"net"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func createServer(t *testing.T) *Server {
+	srv, err := NewServer(context.Background(), ":3000")
+	require.NoError(t, err)
+
+	t.Cleanup(func() {
+		srv.Shutdown()
+	})
+
+	return srv
+}
+
+func createServerWithExistingTopic(t *testing.T, topicName string) *Server {
+	srv := createServer(t)
+	srv.topics[topicName] = topic{
+		name:          topicName,
+		subscriptions: make(map[net.Addr]Subscriber),
+	}
+
+	return srv
+}
+
+func createConnectionAndSubscribe(t *testing.T, topics []string) net.Conn {
+	conn, err := net.Dial("tcp", "localhost:3000")
+	require.NoError(t, err)
+
+	err = binary.Write(conn, binary.BigEndian, Subscribe)
+	require.NoError(t, err)
+
+	rawTopics, err := json.Marshal(topics)
+	require.NoError(t, err)
+
+	err = binary.Write(conn, binary.BigEndian, uint32(len(rawTopics)))
+	require.NoError(t, err)
+
+	_, err = conn.Write(rawTopics)
+	require.NoError(t, err)
+
+	expectedRes := "subscribed"
+
+	buf := make([]byte, len(expectedRes))
+	n, err := conn.Read(buf)
+	require.NoError(t, err)
+	require.Equal(t, len(expectedRes), n)
+
+	assert.Equal(t, expectedRes, string(buf))
+
+	return conn
+}
+
+func TestSubscribeToTopics(t *testing.T) {
+	// create a server with an existing topic so we can test subscribing to a new and
+	// existing topic
+	srv := createServerWithExistingTopic(t, "topic a")
+
+	_ = createConnectionAndSubscribe(t, []string{"topic a", "topic b"})
+
+	assert.Len(t, srv.topics, 2)
+	assert.Len(t, srv.topics["topic a"].subscriptions, 1)
+	assert.Len(t, srv.topics["topic b"].subscriptions, 1)
+}
+
+func TestUnsubscribesFromTopic(t *testing.T) {
+	srv := createServerWithExistingTopic(t, "topic a")
+
+	conn := createConnectionAndSubscribe(t, []string{"topic a", "topic b", "topic c"})
+
+	assert.Len(t, srv.topics, 3)
+	assert.Len(t, srv.topics["topic a"].subscriptions, 1)
+	assert.Len(t, srv.topics["topic b"].subscriptions, 1)
+	assert.Len(t, srv.topics["topic c"].subscriptions, 1)
+
+	err := binary.Write(conn, binary.BigEndian, Unsubscribe)
+	require.NoError(t, err)
+
+	topics := []string{"topic a", "topic b"}
+	rawTopics, err := json.Marshal(topics)
+	require.NoError(t, err)
+
+	err = binary.Write(conn, binary.BigEndian, uint32(len(rawTopics)))
+	require.NoError(t, err)
+
+	_, err = conn.Write(rawTopics)
+	require.NoError(t, err)
+
+	expectedRes := "unsubscribed"
+
+	buf := make([]byte, len(expectedRes))
+	n, err := conn.Read(buf)
+	require.NoError(t, err)
+	require.Equal(t, len(expectedRes), n)
+
+	assert.Equal(t, expectedRes, string(buf))
+
+	assert.Len(t, srv.topics, 3)
+	assert.Len(t, srv.topics["topic a"].subscriptions, 0)
+	assert.Len(t, srv.topics["topic b"].subscriptions, 0)
+	assert.Len(t, srv.topics["topic c"].subscriptions, 1)
+}
+
+func TestSubscriberClosesWithoutUnsubscribing(t *testing.T) {
+	srv := createServer(t)
+
+	conn := createConnectionAndSubscribe(t, []string{"topic a", "topic b"})
+
+	assert.Len(t, srv.topics, 2)
+	assert.Len(t, srv.topics["topic a"].subscriptions, 1)
+	assert.Len(t, srv.topics["topic b"].subscriptions, 1)
+
+	// close the conn
+	err := conn.Close()
+	require.NoError(t, err)
+
+	publisherConn, err := net.Dial("tcp", "localhost:3000")
+	require.NoError(t, err)
+
+	err = binary.Write(publisherConn, binary.BigEndian, Publish)
+	require.NoError(t, err)
+
+	data := []byte("hello world")
+	// send data length first
+	err = binary.Write(publisherConn, binary.BigEndian, uint32(len(data)))
+	require.NoError(t, err)
+	n, err := publisherConn.Write(data)
+	require.NoError(t, err)
+	require.Equal(t, len(data), n)
+
+	assert.Len(t, srv.topics, 2)
+	assert.Len(t, srv.topics["topic a"].subscriptions, 0)
+	assert.Len(t, srv.topics["topic b"].subscriptions, 0)
+}
+
+func TestInvalidAction(t *testing.T) {
+	_ = createServer(t)
+
+	conn, err := net.Dial("tcp", "localhost:3000")
+	require.NoError(t, err)
+
+	err = binary.Write(conn, binary.BigEndian, uint8(99))
+	require.NoError(t, err)
+
+	expectedRes := "unknown action"
+
+	buf := make([]byte, len(expectedRes))
+	n, err := conn.Read(buf)
+	require.NoError(t, err)
+	require.Equal(t, len(expectedRes), n)
+
+	assert.Equal(t, expectedRes, string(buf))
+}
+
+func TestInvalidMessagePublished(t *testing.T) {
+	_ = createServer(t)
+
+	publisherConn, err := net.Dial("tcp", "localhost:3000")
+	require.NoError(t, err)
+
+	err = binary.Write(publisherConn, binary.BigEndian, Publish)
+	require.NoError(t, err)
+
+	// send some data
+	data := []byte("this isn't wrapped in a message type")
+
+	// send data length first
+	err = binary.Write(publisherConn, binary.BigEndian, uint32(len(data)))
+	require.NoError(t, err)
+	n, err := publisherConn.Write(data)
+	require.NoError(t, err)
+	require.Equal(t, len(data), n)
+
+	buf := make([]byte, 15)
+	_, err = publisherConn.Read(buf)
+	require.NoError(t, err)
+	assert.Equal(t, "invalid message", string(buf))
+}
+
+func TestSendsDataToTopicSubscribers(t *testing.T) {
+	_ = createServer(t)
+
+	subscribers := make([]net.Conn, 0, 5)
+	for i := 0; i < 5; i++ {
+		subscriberConn := createConnectionAndSubscribe(t, []string{"topic a", "topic b"})
+
+		subscribers = append(subscribers, subscriberConn)
+	}
+
+	publisherConn, err := net.Dial("tcp", "localhost:3000")
+	require.NoError(t, err)
+
+	err = binary.Write(publisherConn, binary.BigEndian, Publish)
+	require.NoError(t, err)
+
+	// send some data
+	data := []byte("hello world")
+	msg := message{
+		Topic: "topic a",
+		Data:  data,
+	}
+
+	rawMsg, err := json.Marshal(msg)
+	require.NoError(t, err)
+
+	// send data length first
+	err = binary.Write(publisherConn, binary.BigEndian, uint32(len(rawMsg)))
+	require.NoError(t, err)
+	n, err := publisherConn.Write(rawMsg)
+	require.NoError(t, err)
+	require.Equal(t, len(rawMsg), n)
+
+	// check the subsribers got the data
+	for _, conn := range subscribers {
+		buf := make([]byte, len(data))
+		n, err := conn.Read(buf)
+		require.NoError(t, err)
+		require.Equal(t, len(data), n)
+
+		assert.Equal(t, data, buf)
+	}
+}

subscriber.go

···
+package main
+
+import (
+	"fmt"
+	"net"
+)
+
+type Subscriber struct {
+	conn          net.Conn
+	currentOffset int
+}
+
+func (s *Subscriber) SendMessage(data []byte) error {
+	_, err := s.conn.Write(data)
+	if err != nil {
+		return fmt.Errorf("failed to write to connection: %w", err)
+	}
+	return nil
+}

topic.go

···
+package main
+
+import (
+	"log/slog"
+	"net"
+	"sync"
+)
+
+type topic struct {
+	name          string
+	subscriptions map[net.Addr]Subscriber
+	mu            sync.Mutex
+}
+
+func newTopic(name string) topic {
+	return topic{
+		name:          name,
+		subscriptions: make(map[net.Addr]Subscriber),
+	}
+}
+
+func (t *topic) addSubscriber(conn net.Conn) {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	slog.Info("adding subscriber", "conn", conn.LocalAddr())
+	t.subscriptions[conn.LocalAddr()] = Subscriber{conn: conn}
+}
+
+func (t *topic) removeSubscriber(addr net.Addr) {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	slog.Info("removing subscriber", "conn", addr)
+	delete(t.subscriptions, addr)
+}
+
+func (t *topic) sendMessageToSubscribers(msg message) {
+	t.mu.Lock()
+	subscribers := t.subscriptions
+	t.mu.Unlock()
+
+	for addr, subscriber := range subscribers {
+		err := subscriber.SendMessage(msg.Data)
+		if err != nil {
+			slog.Error("failed to send to message", "error", err, "conn", addr)
+			continue
+		}
+	}
+}