Refactor app structure + test improvements (#10)

example/main.go

···
 	"log/slog"
 	"time"
 
-	"github.com/willdot/messagebroker/pubsub"
-	"github.com/willdot/messagebroker/server"
+	"github.com/willdot/messagebroker/client"
+	"github.com/willdot/messagebroker/internal/server"
 )
 
 var consumeOnly *bool
···
 		go sendMessages()
 	}
 
-	sub, err := pubsub.NewSubscriber(":3000")
+	sub, err := client.NewSubscriber(":3000")
 	if err != nil {
 		panic(err)
 	}
···
 }
 
 func sendMessages() {
-	publisher, err := pubsub.NewPublisher("localhost:3000")
+	publisher, err := client.NewPublisher("localhost:3000")
 	if err != nil {
 		panic(err)
 	}
···
 	i := 0
 	for {
 		i++
-		msg := pubsub.NewMessage("topic a", []byte(fmt.Sprintf("message %d", i)))
+		msg := client.NewMessage("topic a", []byte(fmt.Sprintf("message %d", i)))
 
 		err = publisher.PublishMessage(msg)
 		if err != nil {

example/server/main.go

···
 	"syscall"
 	"time"
 
-	"github.com/willdot/messagebroker/server"
+	"github.com/willdot/messagebroker/internal/server"
 )
 
 func main() {

internal/messagestore/memory_store.go

···
+package messagestore
+
+import (
+	"sync"
+
+	"github.com/willdot/messagebroker/internal"
+)
+
+// MemoryStore allows messages to be stored in memory
+type MemoryStore struct {
+	mu         sync.Mutex
+	msgs       map[int]internal.Message
+	nextOffset int
+}
+
+// NewMemoryStore initializes a new in memory store
+func NewMemoryStore() *MemoryStore {
+	return &MemoryStore{
+		msgs: make(map[int]internal.Message),
+	}
+}
+
+// Write will write the provided message to the in memory store
+func (m *MemoryStore) Write(msg internal.Message) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	m.msgs[m.nextOffset] = msg
+
+	m.nextOffset++
+
+	return nil
+}
+
+// ReadFrom will read messages from (and including) the provided offset and pass them to the provided handler
+func (m *MemoryStore) ReadFrom(offset int, handleFunc func(msg internal.Message)) {
+	if offset < 0 || offset >= m.nextOffset {
+		return
+	}
+
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	for i := offset; i < len(m.msgs); i++ {
+		handleFunc(m.msgs[i])
+	}
+}

internal/messge.go

···
+package internal
+
+// Message represents a message that can be sent / received
+type Message struct {
+	Data          []byte
+	DeliveryCount int
+}
+
+// NewMessage intializes a new message
+func NewMessage(data []byte) Message {
+	return Message{Data: data, DeliveryCount: 1}
+}

pubsub/message.go client/message.go

···
-package pubsub
+package client
 
 // Message represents a message that can be published or consumed
 type Message struct {

pubsub/publisher.go client/publisher.go

···
-package pubsub
+package client
 
 import (
 	"encoding/binary"
···
 	"net"
 	"sync"
 
-	"github.com/willdot/messagebroker/server"
+	"github.com/willdot/messagebroker/internal/server"
 )
 
 // Publisher allows messages to be published to a server
···
 		// send topic first
 		topic := fmt.Sprintf("topic:%s", message.Topic)
 
-		topicLenB := make([]byte, 4)
-		binary.BigEndian.PutUint32(topicLenB, uint32(len(topic)))
+		topicLenB := make([]byte, 2)
+		binary.BigEndian.PutUint16(topicLenB, uint16(len(topic)))
 
 		headers := append(topicLenB, []byte(topic)...)
 

pubsub/subscriber.go client/subscriber.go

···
-package pubsub
+package client
 
 import (
 	"context"
···
 	"sync"
 	"time"
 
-	"github.com/willdot/messagebroker/server"
+	"github.com/willdot/messagebroker/internal/server"
 )
 
 type connOpp func(conn net.Conn) error
···
 		return nil
 	}
 
-	var dataLen uint32
+	var dataLen uint16
 	err = binary.Read(conn, binary.BigEndian, &dataLen)
 	if err != nil {
 		return fmt.Errorf("received status %s:", resp)
···
 		return nil
 	}
 
-	var dataLen uint32
+	var dataLen uint16
 	err = binary.Read(conn, binary.BigEndian, &dataLen)
 	if err != nil {
 		return fmt.Errorf("received status %s:", resp)
···
 
 func (s *Subscriber) readMessage(ctx context.Context, msgChan chan *Message) error {
 	op := func(conn net.Conn) error {
-		err := s.conn.SetReadDeadline(time.Now().Add(time.Second))
+		err := s.conn.SetReadDeadline(time.Now().Add(time.Millisecond * 300))
 		if err != nil {
 			return err
 		}
 
-		var topicLen uint64
+		var topicLen uint16
 		err = binary.Read(s.conn, binary.BigEndian, &topicLen)
 		if err != nil {
 			// TODO: check if this is needed elsewhere. I'm not sure where the read deadline resets....

pubsub/subscriber_test.go client/subscriber_test.go

···
-package pubsub
+package client
 
 import (
 	"context"
···
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-
-	"github.com/willdot/messagebroker/server"
+	"github.com/willdot/messagebroker/internal/server"
 )
 
 const (
···
 	err = publisher.PublishMessage(msg)
 	require.NoError(t, err)
 
-	time.Sleep(time.Second)
+	// give the consumer some time to read the messages -- TODO: make better!
+	time.Sleep(time.Millisecond * 300)
 	cancel()
 
 	select {
···
 	}
 
 	// give the consumer some time to read the messages -- TODO: make better!
-	time.Sleep(time.Second)
+	time.Sleep(time.Millisecond * 300)
 	cancel()
 
 	select {

server/message_store.go

···
-package server
-
-import (
-	"sync"
-)
-
-// Memory store allows messages to be stored in memory
-type MemoryStore struct {
-	mu         sync.Mutex
-	msgs       map[int]message
-	nextOffset int
-}
-
-// New memory store initializes a new in memory store
-func NewMemoryStore() *MemoryStore {
-	return &MemoryStore{
-		msgs: make(map[int]message),
-	}
-}
-
-// Write will write the provided message to the in memory store
-func (m *MemoryStore) Write(msg message) error {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-
-	m.msgs[m.nextOffset] = msg
-
-	m.nextOffset++
-
-	return nil
-}
-
-// ReadFrom will read messages from (and including) the provided offset and pass them to the provided handler
-func (m *MemoryStore) ReadFrom(offset int, handleFunc func(msg message)) {
-	if offset < 0 || offset >= m.nextOffset {
-		return
-	}
-
-	m.mu.Lock()
-	defer m.mu.Unlock()
-
-	for i := offset; i < len(m.msgs); i++ {
-		handleFunc(m.msgs[i])
-	}
-}

server/peer/peer.go internal/server/peer.go

···
-package peer
+package server
 
 import (
 	"net"
···
 }
 
 // New returns a new peer.
-func New(conn net.Conn) *Peer {
+func NewPeer(conn net.Conn) *Peer {
 	return &Peer{
 		conn: conn,
 	}

server/server.go internal/server/server.go

···
 	"syscall"
 	"time"
 
-	"github.com/willdot/messagebroker/server/peer"
+	"github.com/willdot/messagebroker/internal"
 )
 
 // Action represents the type of action that a peer requests to do
···
 }
 
 func (s *Server) handleConn(conn net.Conn) {
-	peer := peer.New(conn)
+	peer := NewPeer(conn)
 
 	slog.Info("handling connection", "peer", peer.Addr())
 	defer slog.Info("ending connection", "peer", peer.Addr())
···
 	}
 }
 
-func (s *Server) handleSubscribe(peer *peer.Peer) {
+func (s *Server) handleSubscribe(peer *Peer) {
 	slog.Info("handling subscriber", "peer", peer.Addr())
 	// subscribe the peer to the topic
 	s.subscribePeerToTopic(peer)
 
+	s.waitForPeerAction(peer)
+}
+
+func (s *Server) waitForPeerAction(peer *Peer) {
 	// keep handling the peers connection, getting the action from the peer when it wishes to do something else.
 	// once the peers connection ends, it will be unsubscribed from all topics and returned
 	for {
···
 	}
 }
 
-func (s *Server) subscribePeerToTopic(peer *peer.Peer) {
+func (s *Server) subscribePeerToTopic(peer *Peer) {
 	op := func(conn net.Conn) error {
 		// get the topics the peer wishes to subscribe to
-		dataLen, err := dataLength(conn)
+		dataLen, err := dataLengthUint32(conn)
 		if err != nil {
 			slog.Error(err.Error(), "peer", peer.Addr())
 			writeStatus(Error, "invalid data length of topics provided", conn)
···
 	_ = peer.RunConnOperation(op)
 }
 
-func (s *Server) handleUnsubscribe(peer *peer.Peer) {
+func (s *Server) handleUnsubscribe(peer *Peer) {
 	slog.Info("handling unsubscriber", "peer", peer.Addr())
 	op := func(conn net.Conn) error {
 		// get the topics the peer wishes to unsubscribe from
-		dataLen, err := dataLength(conn)
+		dataLen, err := dataLengthUint32(conn)
 		if err != nil {
 			slog.Error(err.Error(), "peer", peer.Addr())
 			writeStatus(Error, "invalid data length of topics provided", conn)
···
 	_ = peer.RunConnOperation(op)
 }
 
-func (s *Server) handlePublish(peer *peer.Peer) {
+func (s *Server) handlePublish(peer *Peer) {
 	slog.Info("handling publisher", "peer", peer.Addr())
 	for {
 		op := func(conn net.Conn) error {
-			dataLen, err := dataLength(conn)
+			topicDataLen, err := dataLengthUint16(conn)
 			if err != nil {
 				if errors.Is(err, io.EOF) {
 					return nil
···
 				writeStatus(Error, "invalid data length of data provided", conn)
 				return nil
 			}
-			if dataLen == 0 {
+			if topicDataLen == 0 {
 				return nil
 			}
-			topicBuf := make([]byte, dataLen)
+			topicBuf := make([]byte, topicDataLen)
 			_, err = conn.Read(topicBuf)
 			if err != nil {
 				slog.Error("failed to read topic from peer", "error", err, "peer", peer.Addr())
···
 			}
 			topicStr = strings.TrimPrefix(topicStr, "topic:")
 
-			dataLen, err = dataLength(conn)
+			msgDataLen, err := dataLengthUint32(conn)
 			if err != nil {
 				slog.Error(err.Error(), "peer", peer.Addr())
 				writeStatus(Error, "invalid data length of data provided", conn)
 				return nil
 			}
-			if dataLen == 0 {
+			if msgDataLen == 0 {
 				return nil
 			}
 
-			dataBuf := make([]byte, dataLen)
+			dataBuf := make([]byte, msgDataLen)
 			_, err = conn.Read(dataBuf)
 			if err != nil {
 				slog.Error("failed to read data from peer", "error", err, "peer", peer.Addr())
···
 				s.topics[topicStr] = topic
 			}
 
-			message := newMessage(dataBuf)
+			message := internal.NewMessage(dataBuf)
 
 			err = topic.sendMessageToSubscribers(message)
 			if err != nil {
···
 	}
 }
 
-func (s *Server) subscribeToTopics(peer *peer.Peer, topics []string, startAt int) {
+func (s *Server) subscribeToTopics(peer *Peer, topics []string, startAt int) {
 	slog.Info("subscribing peer to topics", "topics", topics, "peer", peer.Addr())
 	for _, topic := range topics {
 		s.addSubsciberToTopic(topic, peer, startAt)
 	}
 }
 
-func (s *Server) addSubsciberToTopic(topicName string, peer *peer.Peer, startAt int) {
+func (s *Server) addSubsciberToTopic(topicName string, peer *Peer, startAt int) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
···
 	s.topics[topicName] = t
 }
 
-func (s *Server) unsubscribeToTopics(peer *peer.Peer, topics []string) {
+func (s *Server) unsubscribeToTopics(peer *Peer, topics []string) {
 	slog.Info("unsubscribing peer from topics", "topics", topics, "peer", peer.Addr())
 	for _, topic := range topics {
 		s.removeSubsciberFromTopic(topic, peer)
 	}
 }
 
-func (s *Server) removeSubsciberFromTopic(topicName string, peer *peer.Peer) {
+func (s *Server) removeSubsciberFromTopic(topicName string, peer *Peer) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
···
 	delete(t.subscriptions, peer.Addr())
 }
 
-func (s *Server) unsubscribePeerFromAllTopics(peer *peer.Peer) {
+func (s *Server) unsubscribePeerFromAllTopics(peer *Peer) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
···
 	return nil
 }
 
-func readAction(peer *peer.Peer, timeout time.Duration) (Action, error) {
+func readAction(peer *Peer, timeout time.Duration) (Action, error) {
 	var action Action
 	op := func(conn net.Conn) error {
 		if timeout > 0 {
···
 	return action, nil
 }
 
-func writeInvalidAction(peer *peer.Peer) {
+func writeInvalidAction(peer *Peer) {
 	op := func(conn net.Conn) error {
 		writeStatus(Error, "unknown action", conn)
 		return nil
···
 	_ = peer.RunConnOperation(op)
 }
 
-func dataLength(conn net.Conn) (uint32, error) {
+func dataLengthUint32(conn net.Conn) (uint32, error) {
 	var dataLen uint32
+	err := binary.Read(conn, binary.BigEndian, &dataLen)
+	if err != nil {
+		return 0, err
+	}
+	return dataLen, nil
+}
+
+func dataLengthUint16(conn net.Conn) (uint16, error) {
+	var dataLen uint16
 	err := binary.Read(conn, binary.BigEndian, &dataLen)
 	if err != nil {
 		return 0, err
···
 	headers := statusB
 
 	if len(message) > 0 {
-		sizeB := make([]byte, 4)
-		binary.BigEndian.PutUint32(sizeB, uint32(len(message)))
+		sizeB := make([]byte, 2)
+		binary.BigEndian.PutUint16(sizeB, uint16(len(message)))
 		headers = append(headers, sizeB...)
 	}
 

server/server_test.go internal/server/server_test.go

···
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"github.com/willdot/messagebroker/internal/messagestore"
 )
 
 const (
···
 	srv.topics[topicName] = &topic{
 		name:          topicName,
 		subscriptions: make(map[net.Addr]*subscriber),
-		messageStore:  NewMemoryStore(),
+		messageStore:  messagestore.NewMemoryStore(),
 	}
 
 	return srv
···
 
 	expectedMessage := "unknown action"
 
-	var dataLen uint32
+	var dataLen uint16
 	err = binary.Read(conn, binary.BigEndian, &dataLen)
 	require.NoError(t, err)
 	assert.Equal(t, len(expectedMessage), int(dataLen))
···
 
 	expectedMessage := "topic data does not contain 'topic:' prefix"
 
-	var dataLen uint32
+	var dataLen uint16
 	err = binary.Read(publisherConn, binary.BigEndian, &dataLen)
 	require.NoError(t, err)
 	assert.Equal(t, len(expectedMessage), int(dataLen))
···
 
 	// check the subsribers got the data
 	readMessage := func(conn net.Conn, ack Action) {
-		var topicLen uint64
+		var topicLen uint16
 		err = binary.Read(conn, binary.BigEndian, &topicLen)
 		require.NoError(t, err)
 
···
 	readMessage(subscriberConn, Ack)
 	// reading for another message should now timeout but give enough time for the ack delay to kick in
 	// should the second read of the message not have been ack'd properly
-	var topicLen uint64
+	var topicLen uint16
 	_ = subscriberConn.SetReadDeadline(time.Now().Add(ackDelay + time.Millisecond*100))
 	err = binary.Read(subscriberConn, binary.BigEndian, &topicLen)
 	require.Error(t, err)
···
 
 	// check the subsribers got the data
 	readMessage := func(conn net.Conn, ack bool) {
-		var topicLen uint64
+		var topicLen uint16
 		err = binary.Read(conn, binary.BigEndian, &topicLen)
 		require.NoError(t, err)
 
···
 
 	// reading for another message should now timeout but give enough time for the ack delay to kick in
 	// should the second read of the message not have been ack'd properly
-	var topicLen uint64
+	var topicLen uint16
 	_ = subscriberConn.SetReadDeadline(time.Now().Add(ackDelay + time.Millisecond*100))
 	err = binary.Read(subscriberConn, binary.BigEndian, &topicLen)
 	require.Error(t, err)
···
 
 	// check the subsribers got the data
 	readMessage := func(conn net.Conn, ack bool) {
-		var topicLen uint64
+		var topicLen uint16
 		err = binary.Read(conn, binary.BigEndian, &topicLen)
 		require.NoError(t, err)
 
···
 	readMessage(subscriberConn, false)
 
 	// reading for the message should now timeout as we have nack'd the message too many times
-	var topicLen uint64
+	var topicLen uint16
 	_ = subscriberConn.SetReadDeadline(time.Now().Add(ackDelay + time.Millisecond*100))
 	err = binary.Read(subscriberConn, binary.BigEndian, &topicLen)
 	require.Error(t, err)
···
 }
 
 func readMessage(t *testing.T, subscriberConn net.Conn) []byte {
-	var topicLen uint64
+	var topicLen uint16
 	err := binary.Read(subscriberConn, binary.BigEndian, &topicLen)
 	require.NoError(t, err)
 

server/subscriber.go internal/server/subscriber.go

···
 	"net"
 	"time"
 
-	"github.com/willdot/messagebroker/server/peer"
+	"github.com/willdot/messagebroker/internal"
 )
 
 type subscriber struct {
-	peer          *peer.Peer
+	peer          *Peer
 	topic         string
-	messages      chan message
+	messages      chan internal.Message
 	unsubscribeCh chan struct{}
 
 	ackDelay   time.Duration
 	ackTimeout time.Duration
 }
 
-type message struct {
-	data          []byte
-	deliveryCount int
-}
-
-func newMessage(data []byte) message {
-	return message{data: data, deliveryCount: 1}
-}
-
-func newSubscriber(peer *peer.Peer, topic *topic, ackDelay, ackTimeout time.Duration, startAt int) *subscriber {
+func newSubscriber(peer *Peer, topic *topic, ackDelay, ackTimeout time.Duration, startAt int) *subscriber {
 	s := &subscriber{
 		peer:          peer,
 		topic:         topic.name,
-		messages:      make(chan message),
+		messages:      make(chan internal.Message),
 		ackDelay:      ackDelay,
 		ackTimeout:    ackTimeout,
 		unsubscribeCh: make(chan struct{}, 1),
···
 	go s.sendMessages()
 
 	go func() {
-		topic.messageStore.ReadFrom(startAt, func(msg message) {
+		topic.messageStore.ReadFrom(startAt, func(msg internal.Message) {
 			select {
 			case s.messages <- msg:
 				return
···
 				continue
 			}
 
-			if msg.deliveryCount >= 5 {
+			if msg.DeliveryCount >= 5 {
 				slog.Error("max delivery count for message. Dropping", "peer", s.peer.Addr())
 				continue
 			}
 
-			msg.deliveryCount++
+			msg.DeliveryCount++
 			s.addMessage(msg, s.ackDelay)
 		}
 	}
 }
 
-func (s *subscriber) addMessage(msg message, delay time.Duration) {
+func (s *subscriber) addMessage(msg internal.Message, delay time.Duration) {
 	go func() {
 		timer := time.NewTimer(delay)
 		defer timer.Stop()
···
 	}()
 }
 
-func (s *subscriber) sendMessage(topic string, msg message) (bool, error) {
+func (s *subscriber) sendMessage(topic string, msg internal.Message) (bool, error) {
 	var ack bool
 	op := func(conn net.Conn) error {
-		// TODO: why did I chose uint64 for topic len?
-		topicB := make([]byte, 8)
-		binary.BigEndian.PutUint64(topicB, uint64(len(topic)))
+		topicB := make([]byte, 2)
+		binary.BigEndian.PutUint16(topicB, uint16(len(topic)))
 
 		headers := topicB
 		headers = append(headers, []byte(topic)...)
 
 		// TODO: if message is empty, return error?
 		dataLenB := make([]byte, 8)
-		binary.BigEndian.PutUint64(dataLenB, uint64(len(msg.data)))
+		binary.BigEndian.PutUint64(dataLenB, uint64(len(msg.Data)))
 		headers = append(headers, dataLenB...)
 
-		_, err := conn.Write(append(headers, msg.data...))
+		_, err := conn.Write(append(headers, msg.Data...))
 		if err != nil {
 			return fmt.Errorf("failed to write to peer: %w", err)
 		}
 
-		var ackRes Action
 		if err := conn.SetReadDeadline(time.Now().Add(s.ackTimeout)); err != nil {
 			slog.Error("failed to set connection read deadline", "error", err, "peer", s.peer.Addr())
 		}
···
 				slog.Error("failed to reset connection read deadline", "error", err, "peer", s.peer.Addr())
 			}
 		}()
+		var ackRes Action
 		err = binary.Read(conn, binary.BigEndian, &ackRes)
 		if err != nil {
 			return fmt.Errorf("failed to read ack from peer: %w", err)

server/topic.go internal/server/topic.go

···
 	"fmt"
 	"net"
 	"sync"
+
+	"github.com/willdot/messagebroker/internal"
+	"github.com/willdot/messagebroker/internal/messagestore"
 )
 
 type Store interface {
-	Write(msg message) error
-	ReadFrom(offset int, handleFunc func(msg message))
+	Write(msg internal.Message) error
+	ReadFrom(offset int, handleFunc func(msg internal.Message))
 }
 
 type topic struct {
···
 }
 
 func newTopic(name string) *topic {
-	messageStore := NewMemoryStore()
+	messageStore := messagestore.NewMemoryStore()
 	return &topic{
 		name:          name,
 		subscriptions: make(map[net.Addr]*subscriber),
···
 	}
 }
 
-func (t *topic) sendMessageToSubscribers(msg message) error {
+func (t *topic) sendMessageToSubscribers(msg internal.Message) error {
 	err := t.messageStore.Write(msg)
 	if err != nil {
 		return fmt.Errorf("failed to write message to store: %w", err)