An experimental pub/sub client and server project.

Refactor. Huge refactor to make conns synchronous

+2 -1
.gitignore
··· 1 - .DS_STORE
··· 1 + .DS_STORE 2 + example/example
+20
dockerfile.example-server
···
··· 1 + FROM golang:latest as builder 2 + 3 + WORKDIR /app 4 + 5 + COPY go.mod go.sum ./ 6 + COPY example/server/ ./ 7 + RUN go mod download 8 + 9 + COPY . . 10 + 11 + RUN CGO_ENABLED=0 go build -o message-broker-server . 12 + 13 + FROM alpine:latest 14 + 15 + RUN apk --no-cache add ca-certificates 16 + 17 + WORKDIR /root/ 18 + COPY --from=builder /app/message-broker-server . 19 + 20 + CMD ["./message-broker-server"]
+9 -9
example/main.go
··· 2 3 import ( 4 "context" 5 "fmt" 6 "log/slog" 7 8 - "github.com/willdot/messagebroker" 9 "github.com/willdot/messagebroker/pubsub" 10 - "github.com/willdot/messagebroker/server" 11 ) 12 13 func main() { 14 - server, err := server.New(context.Background(), ":3000") 15 - if err != nil { 16 - panic(err) 17 - } 18 - defer server.Shutdown() 19 20 - go sendMessages() 21 22 sub, err := pubsub.NewSubscriber(":3000") 23 if err != nil { ··· 49 i := 0 50 for { 51 i++ 52 - msg := messagebroker.Message{ 53 Topic: "topic a", 54 Data: []byte(fmt.Sprintf("message %d", i)), 55 }
··· 2 3 import ( 4 "context" 5 + "flag" 6 "fmt" 7 "log/slog" 8 9 "github.com/willdot/messagebroker/pubsub" 10 ) 11 12 + var consumeOnly *bool 13 + 14 func main() { 15 + consumeOnly = flag.Bool("consume-only", false, "just consumes (doesn't start server and doesn't publish)") 16 + flag.Parse() 17 18 + if *consumeOnly == false { 19 + go sendMessages() 20 + } 21 22 sub, err := pubsub.NewSubscriber(":3000") 23 if err != nil { ··· 49 i := 0 50 for { 51 i++ 52 + msg := pubsub.Message{ 53 Topic: "topic a", 54 Data: []byte(fmt.Sprintf("message %d", i)), 55 }
+23
example/server/main.go
···
··· 1 + package main 2 + 3 + import ( 4 + "log" 5 + "os" 6 + "os/signal" 7 + "syscall" 8 + 9 + "github.com/willdot/messagebroker/server" 10 + ) 11 + 12 + func main() { 13 + srv, err := server.New(":3000") 14 + if err != nil { 15 + log.Fatal(err) 16 + } 17 + defer srv.Shutdown() 18 + 19 + signals := make(chan os.Signal, 1) 20 + signal.Notify(signals, syscall.SIGTERM, syscall.SIGINT) 21 + 22 + <-signals 23 + }
+5 -1
go.mod
··· 2 3 go 1.21.0 4 5 - require github.com/stretchr/testify v1.8.4 6 7 require ( 8 github.com/davecgh/go-spew v1.1.1 // indirect
··· 2 3 go 1.21.0 4 5 + require ( 6 + github.com/docker/distribution v2.8.3+incompatible 7 + github.com/google/uuid v1.4.0 8 + github.com/stretchr/testify v1.8.4 9 + ) 10 11 require ( 12 github.com/davecgh/go-spew v1.1.1 // indirect
+4
go.sum
··· 1 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= 2 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 3 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 4 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 5 github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
··· 1 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= 2 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 3 + github.com/docker/distribution v2.8.3+incompatible h1:AtKxIZ36LoNK51+Z6RpzLpddBirtxJnzDrHLEKxTAYk= 4 + github.com/docker/distribution v2.8.3+incompatible/go.mod h1:J2gT2udsDAN96Uj4KfcMRqY0/ypR+oyYUYmja8H+y+w= 5 + github.com/google/uuid v1.4.0 h1:MtMxsa51/r9yyhkyLsVeVt0B+BGQZzpQiTQ4eHZ8bc4= 6 + github.com/google/uuid v1.4.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= 7 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 8 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 9 github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
+1 -1
message.go pubsub/message.go
··· 1 - package messagebroker 2 3 // Message represents a message that can be published or consumed 4 type Message struct {
··· 1 + package pubsub 2 3 // Message represents a message that can be published or consumed 4 type Message struct {
+33 -16
pubsub/publisher.go
··· 2 3 import ( 4 "encoding/binary" 5 - "encoding/json" 6 "fmt" 7 "net" 8 9 - "github.com/willdot/messagebroker" 10 "github.com/willdot/messagebroker/server" 11 ) 12 13 // Publisher allows messages to be published to a server 14 type Publisher struct { 15 - conn net.Conn 16 } 17 18 // NewPublisher connects to the server at the given address and registers as a publisher ··· 39 } 40 41 // Publish will publish the given message to the server 42 - func (p *Publisher) PublishMessage(message messagebroker.Message) error { 43 - b, err := json.Marshal(message) 44 - if err != nil { 45 - return fmt.Errorf("failed to marshal message: %w", err) 46 } 47 48 - err = binary.Write(p.conn, binary.BigEndian, uint32(len(b))) 49 - if err != nil { 50 - return fmt.Errorf("failed to write message size to server") 51 - } 52 53 - _, err = p.conn.Write(b) 54 - if err != nil { 55 - return fmt.Errorf("failed to publish data to server") 56 - } 57 58 - return nil 59 }
··· 2 3 import ( 4 "encoding/binary" 5 "fmt" 6 "net" 7 + "sync" 8 9 "github.com/willdot/messagebroker/server" 10 ) 11 12 // Publisher allows messages to be published to a server 13 type Publisher struct { 14 + conn net.Conn 15 + connMu sync.Mutex 16 } 17 18 // NewPublisher connects to the server at the given address and registers as a publisher ··· 39 } 40 41 // Publish will publish the given message to the server 42 + func (p *Publisher) PublishMessage(message Message) error { 43 + op := func(conn net.Conn) error { 44 + // send topic first 45 + topic := fmt.Sprintf("topic:%s", message.Topic) 46 + err := binary.Write(p.conn, binary.BigEndian, uint32(len(topic))) 47 + if err != nil { 48 + return fmt.Errorf("failed to write topic size to server") 49 + } 50 + 51 + _, err = p.conn.Write([]byte(topic)) 52 + if err != nil { 53 + return fmt.Errorf("failed to write topic to server") 54 + } 55 + 56 + err = binary.Write(p.conn, binary.BigEndian, uint32(len(message.Data))) 57 + if err != nil { 58 + return fmt.Errorf("failed to write message size to server") 59 + } 60 + 61 + _, err = p.conn.Write(message.Data) 62 + if err != nil { 63 + return fmt.Errorf("failed to publish data to server") 64 + } 65 + return nil 66 } 67 68 + return p.connOperation(op) 69 + } 70 71 + func (p *Publisher) connOperation(op connOpp) error { 72 + p.connMu.Lock() 73 + defer p.connMu.Unlock() 74 75 + return op(p.conn) 76 }
+141 -97
pubsub/subscriber.go
··· 4 "context" 5 "encoding/binary" 6 "encoding/json" 7 "fmt" 8 - "log/slog" 9 "net" 10 "time" 11 12 - "github.com/willdot/messagebroker" 13 "github.com/willdot/messagebroker/server" 14 ) 15 16 // Subscriber allows subscriptions to a server and the consumption of messages 17 type Subscriber struct { 18 - conn net.Conn 19 } 20 21 // NewSubscriber will connect to the server at the given address ··· 37 38 // SubscribeToTopics will subscribe to the provided topics 39 func (s *Subscriber) SubscribeToTopics(topicNames []string) error { 40 - err := binary.Write(s.conn, binary.BigEndian, server.Subscribe) 41 - if err != nil { 42 - return fmt.Errorf("failed to subscribe: %w", err) 43 - } 44 45 - b, err := json.Marshal(topicNames) 46 - if err != nil { 47 - return fmt.Errorf("failed to marshal topic names: %w", err) 48 - } 49 50 - err = binary.Write(s.conn, binary.BigEndian, uint32(len(b))) 51 - if err != nil { 52 - return fmt.Errorf("failed to write topic data length: %w", err) 53 - } 54 55 - _, err = s.conn.Write(b) 56 - if err != nil { 57 - return fmt.Errorf("failed to subscribe to topics: %w", err) 58 - } 59 60 - var resp server.Status 61 - err = binary.Read(s.conn, binary.BigEndian, &resp) 62 - if err != nil { 63 - return fmt.Errorf("failed to read confirmation of subscription: %w", err) 64 - } 65 66 - if resp == server.Subscribed { 67 - return nil 68 - } 69 70 - var dataLen uint32 71 - err = binary.Read(s.conn, binary.BigEndian, &dataLen) 72 - if err != nil { 73 - return fmt.Errorf("received status %s:", resp) 74 - } 75 76 - buf := make([]byte, dataLen) 77 - _, err = s.conn.Read(buf) 78 - if err != nil { 79 - return fmt.Errorf("received status %s:", resp) 80 } 81 82 - return fmt.Errorf("received status %s - %s", resp, buf) 83 } 84 85 // UnsubscribeToTopics will unsubscribe to the provided topics 86 func (s *Subscriber) UnsubscribeToTopics(topicNames []string) error { 87 - err := binary.Write(s.conn, binary.BigEndian, server.Unsubscribe) 88 - if err != nil { 89 - return fmt.Errorf("failed to unsubscribe: %w", err) 90 - } 91 92 - b, err := json.Marshal(topicNames) 93 - if err != nil { 94 - return fmt.Errorf("failed to marshal topic names: %w", err) 95 - } 96 97 - err = binary.Write(s.conn, binary.BigEndian, uint32(len(b))) 98 - if err != nil { 99 - return fmt.Errorf("failed to write topic data length: %w", err) 100 - } 101 102 - _, err = s.conn.Write(b) 103 - if err != nil { 104 - return fmt.Errorf("failed to unsubscribe to topics: %w", err) 105 - } 106 107 - var resp server.Status 108 - err = binary.Read(s.conn, binary.BigEndian, &resp) 109 - if err != nil { 110 - return fmt.Errorf("failed to read confirmation of unsubscription: %w", err) 111 - } 112 113 - if resp == server.Unsubscribed { 114 - return nil 115 - } 116 117 - var dataLen uint32 118 - err = binary.Read(s.conn, binary.BigEndian, &dataLen) 119 - if err != nil { 120 - return fmt.Errorf("received status %s:", resp) 121 - } 122 123 - buf := make([]byte, dataLen) 124 - _, err = s.conn.Read(buf) 125 - if err != nil { 126 - return fmt.Errorf("received status %s:", resp) 127 } 128 129 - return fmt.Errorf("received status %s - %s", resp, buf) 130 } 131 132 // Consumer allows the consumption of messages. If during the consumer receiving messages from the 133 // server an error occurs, it will be stored in Err 134 type Consumer struct { 135 - msgs chan messagebroker.Message 136 // TODO: better error handling? Maybe a channel of errors? 137 Err error 138 } 139 140 // Messages returns a channel in which this consumer will put messages onto. It is safe to range over the channel since it will be closed once 141 // the consumer has finished either due to an error or from being cancelled. 142 - func (c *Consumer) Messages() <-chan messagebroker.Message { 143 return c.msgs 144 } 145 ··· 147 // to read the messages 148 func (s *Subscriber) Consume(ctx context.Context) *Consumer { 149 consumer := &Consumer{ 150 - msgs: make(chan messagebroker.Message), 151 } 152 153 go s.consume(ctx, consumer) ··· 174 } 175 } 176 177 - func (s *Subscriber) readMessage() (*messagebroker.Message, error) { 178 - err := s.conn.SetReadDeadline(time.Now().Add(time.Second)) 179 - if err != nil { 180 - return nil, err 181 } 182 183 - var dataLen uint64 184 - err = binary.Read(s.conn, binary.BigEndian, &dataLen) 185 if err != nil { 186 - if neterr, ok := err.(net.Error); ok && neterr.Timeout() { 187 return nil, nil 188 } 189 return nil, err 190 } 191 192 - if dataLen <= 0 { 193 - return nil, nil 194 - } 195 196 - buf := make([]byte, dataLen) 197 - _, err = s.conn.Read(buf) 198 - if err != nil { 199 - return nil, err 200 - } 201 202 - var msg messagebroker.Message 203 - err = json.Unmarshal(buf, &msg) 204 - if err != nil { 205 - slog.Error("failed to unmarshal message", "error", err) 206 - return nil, nil 207 - } 208 - 209 - return &msg, nil 210 }
··· 4 "context" 5 "encoding/binary" 6 "encoding/json" 7 + "errors" 8 "fmt" 9 "net" 10 + "sync" 11 "time" 12 13 "github.com/willdot/messagebroker/server" 14 ) 15 16 + type connOpp func(conn net.Conn) error 17 + 18 // Subscriber allows subscriptions to a server and the consumption of messages 19 type Subscriber struct { 20 + conn net.Conn 21 + connMu sync.Mutex 22 } 23 24 // NewSubscriber will connect to the server at the given address ··· 40 41 // SubscribeToTopics will subscribe to the provided topics 42 func (s *Subscriber) SubscribeToTopics(topicNames []string) error { 43 + op := func(conn net.Conn) error { 44 + err := binary.Write(conn, binary.BigEndian, server.Subscribe) 45 + if err != nil { 46 + return fmt.Errorf("failed to subscribe: %w", err) 47 + } 48 49 + b, err := json.Marshal(topicNames) 50 + if err != nil { 51 + return fmt.Errorf("failed to marshal topic names: %w", err) 52 + } 53 54 + err = binary.Write(conn, binary.BigEndian, uint32(len(b))) 55 + if err != nil { 56 + return fmt.Errorf("failed to write topic data length: %w", err) 57 + } 58 59 + _, err = conn.Write(b) 60 + if err != nil { 61 + return fmt.Errorf("failed to subscribe to topics: %w", err) 62 + } 63 64 + var resp server.Status 65 + err = binary.Read(conn, binary.BigEndian, &resp) 66 + if err != nil { 67 + return fmt.Errorf("failed to read confirmation of subscription: %w", err) 68 + } 69 + 70 + if resp == server.Subscribed { 71 + return nil 72 + } 73 74 + var dataLen uint32 75 + err = binary.Read(conn, binary.BigEndian, &dataLen) 76 + if err != nil { 77 + return fmt.Errorf("received status %s:", resp) 78 + } 79 80 + buf := make([]byte, dataLen) 81 + _, err = conn.Read(buf) 82 + if err != nil { 83 + return fmt.Errorf("received status %s:", resp) 84 + } 85 86 + return fmt.Errorf("received status %s - %s", resp, buf) 87 } 88 89 + return s.connOperation(op) 90 } 91 92 // UnsubscribeToTopics will unsubscribe to the provided topics 93 func (s *Subscriber) UnsubscribeToTopics(topicNames []string) error { 94 + op := func(conn net.Conn) error { 95 + err := binary.Write(conn, binary.BigEndian, server.Unsubscribe) 96 + if err != nil { 97 + return fmt.Errorf("failed to unsubscribe: %w", err) 98 + } 99 100 + b, err := json.Marshal(topicNames) 101 + if err != nil { 102 + return fmt.Errorf("failed to marshal topic names: %w", err) 103 + } 104 105 + err = binary.Write(conn, binary.BigEndian, uint32(len(b))) 106 + if err != nil { 107 + return fmt.Errorf("failed to write topic data length: %w", err) 108 + } 109 110 + _, err = conn.Write(b) 111 + if err != nil { 112 + return fmt.Errorf("failed to unsubscribe to topics: %w", err) 113 + } 114 115 + var resp server.Status 116 + err = binary.Read(conn, binary.BigEndian, &resp) 117 + if err != nil { 118 + return fmt.Errorf("failed to read confirmation of unsubscription: %w", err) 119 + } 120 121 + if resp == server.Unsubscribed { 122 + return nil 123 + } 124 125 + var dataLen uint32 126 + err = binary.Read(conn, binary.BigEndian, &dataLen) 127 + if err != nil { 128 + return fmt.Errorf("received status %s:", resp) 129 + } 130 131 + buf := make([]byte, dataLen) 132 + _, err = conn.Read(buf) 133 + if err != nil { 134 + return fmt.Errorf("received status %s:", resp) 135 + } 136 + 137 + return fmt.Errorf("received status %s - %s", resp, buf) 138 } 139 140 + return s.connOperation(op) 141 } 142 143 // Consumer allows the consumption of messages. If during the consumer receiving messages from the 144 // server an error occurs, it will be stored in Err 145 type Consumer struct { 146 + msgs chan Message 147 // TODO: better error handling? Maybe a channel of errors? 148 Err error 149 } 150 151 // Messages returns a channel in which this consumer will put messages onto. It is safe to range over the channel since it will be closed once 152 // the consumer has finished either due to an error or from being cancelled. 153 + func (c *Consumer) Messages() <-chan Message { 154 return c.msgs 155 } 156 ··· 158 // to read the messages 159 func (s *Subscriber) Consume(ctx context.Context) *Consumer { 160 consumer := &Consumer{ 161 + msgs: make(chan Message), 162 } 163 164 go s.consume(ctx, consumer) ··· 185 } 186 } 187 188 + func (s *Subscriber) readMessage() (*Message, error) { 189 + var msg *Message 190 + op := func(conn net.Conn) error { 191 + err := s.conn.SetReadDeadline(time.Now().Add(time.Second)) 192 + if err != nil { 193 + return err 194 + } 195 + 196 + var topicLen uint64 197 + err = binary.Read(s.conn, binary.BigEndian, &topicLen) 198 + if err != nil { 199 + // TODO: check if this is needed elsewhere. I'm not sure where the read deadline resets.... 200 + if neterr, ok := err.(net.Error); ok && neterr.Timeout() { 201 + return nil 202 + } 203 + return err 204 + } 205 + 206 + topicBuf := make([]byte, topicLen) 207 + _, err = s.conn.Read(topicBuf) 208 + if err != nil { 209 + return err 210 + } 211 + 212 + var dataLen uint64 213 + err = binary.Read(s.conn, binary.BigEndian, &dataLen) 214 + if err != nil { 215 + return err 216 + } 217 + 218 + if dataLen <= 0 { 219 + return nil 220 + } 221 + 222 + dataBuf := make([]byte, dataLen) 223 + _, err = s.conn.Read(dataBuf) 224 + if err != nil { 225 + return err 226 + } 227 + 228 + msg = &Message{ 229 + Data: dataBuf, 230 + Topic: string(topicBuf), 231 + } 232 + 233 + return nil 234 + 235 } 236 237 + err := s.connOperation(op) 238 if err != nil { 239 + var neterr net.Error 240 + if errors.As(err, &neterr) && neterr.Timeout() { 241 return nil, nil 242 } 243 return nil, err 244 } 245 246 + return msg, err 247 + } 248 249 + func (s *Subscriber) connOperation(op connOpp) error { 250 + s.connMu.Lock() 251 + defer s.connMu.Unlock() 252 253 + return op(s.conn) 254 }
+19 -22
pubsub/subscriber_test.go
··· 8 9 "github.com/stretchr/testify/assert" 10 "github.com/stretchr/testify/require" 11 - "github.com/willdot/messagebroker" 12 13 "github.com/willdot/messagebroker/server" 14 ) 15 16 const ( 17 - serverAddr = ":3000" 18 ) 19 20 func createServer(t *testing.T) { 21 - server, err := server.New(context.Background(), serverAddr) 22 require.NoError(t, err) 23 24 t.Cleanup(func() { ··· 72 sub.Close() 73 }) 74 75 - topics := []string{"topic a", "topic b"} 76 77 err = sub.SubscribeToTopics(topics) 78 require.NoError(t, err) ··· 88 sub.Close() 89 }) 90 91 - topics := []string{"topic a", "topic b"} 92 93 err = sub.SubscribeToTopics(topics) 94 require.NoError(t, err) 95 96 - err = sub.UnsubscribeToTopics([]string{"topic a"}) 97 require.NoError(t, err) 98 99 ctx, cancel := context.WithCancel(context.Background()) ··· 104 consumer := sub.Consume(ctx) 105 require.NoError(t, err) 106 107 - var receivedMessages []messagebroker.Message 108 consumerFinCh := make(chan struct{}) 109 go func() { 110 for msg := range consumer.Messages() { ··· 118 // publish a message to both topics and check the subscriber only gets the message from the 1 topic 119 // and not the unsubscribed topic 120 121 - publisher, err := NewPublisher("localhost:3000") 122 require.NoError(t, err) 123 t.Cleanup(func() { 124 publisher.Close() 125 }) 126 127 - msg := messagebroker.Message{ 128 - Topic: "topic a", 129 Data: []byte("hello world"), 130 } 131 132 err = publisher.PublishMessage(msg) 133 require.NoError(t, err) 134 135 - msg.Topic = "topic b" 136 err = publisher.PublishMessage(msg) 137 require.NoError(t, err) 138 139 cancel() 140 141 - // give the consumer some time to read the messages -- TODO: make better! 142 - time.Sleep(time.Millisecond * 500) 143 - cancel() 144 - 145 select { 146 case <-consumerFinCh: 147 break ··· 150 } 151 152 assert.Len(t, receivedMessages, 1) 153 - assert.Equal(t, "topic b", receivedMessages[0].Topic) 154 } 155 156 func TestPublishAndSubscribe(t *testing.T) { ··· 163 sub.Close() 164 }) 165 166 - topics := []string{"topic a", "topic b"} 167 168 err = sub.SubscribeToTopics(topics) 169 require.NoError(t, err) ··· 176 consumer := sub.Consume(ctx) 177 require.NoError(t, err) 178 179 - var receivedMessages []messagebroker.Message 180 181 consumerFinCh := make(chan struct{}) 182 go func() { ··· 188 consumerFinCh <- struct{}{} 189 }() 190 191 - publisher, err := NewPublisher("localhost:3000") 192 require.NoError(t, err) 193 t.Cleanup(func() { 194 publisher.Close() 195 }) 196 197 // send some messages 198 - sentMessages := make([]messagebroker.Message, 0, 10) 199 for i := 0; i < 10; i++ { 200 - msg := messagebroker.Message{ 201 - Topic: "topic a", 202 Data: []byte(fmt.Sprintf("message %d", i)), 203 } 204
··· 8 9 "github.com/stretchr/testify/assert" 10 "github.com/stretchr/testify/require" 11 12 "github.com/willdot/messagebroker/server" 13 ) 14 15 const ( 16 + serverAddr = ":9999" 17 + topicA = "topic a" 18 + topicB = "topic b" 19 ) 20 21 func createServer(t *testing.T) { 22 + server, err := server.New(serverAddr) 23 require.NoError(t, err) 24 25 t.Cleanup(func() { ··· 73 sub.Close() 74 }) 75 76 + topics := []string{topicA, topicB} 77 78 err = sub.SubscribeToTopics(topics) 79 require.NoError(t, err) ··· 89 sub.Close() 90 }) 91 92 + topics := []string{topicA, topicB} 93 94 err = sub.SubscribeToTopics(topics) 95 require.NoError(t, err) 96 97 + err = sub.UnsubscribeToTopics([]string{topicA}) 98 require.NoError(t, err) 99 100 ctx, cancel := context.WithCancel(context.Background()) ··· 105 consumer := sub.Consume(ctx) 106 require.NoError(t, err) 107 108 + var receivedMessages []Message 109 consumerFinCh := make(chan struct{}) 110 go func() { 111 for msg := range consumer.Messages() { ··· 119 // publish a message to both topics and check the subscriber only gets the message from the 1 topic 120 // and not the unsubscribed topic 121 122 + publisher, err := NewPublisher("localhost:9999") 123 require.NoError(t, err) 124 t.Cleanup(func() { 125 publisher.Close() 126 }) 127 128 + msg := Message{ 129 + Topic: topicA, 130 Data: []byte("hello world"), 131 } 132 133 err = publisher.PublishMessage(msg) 134 require.NoError(t, err) 135 136 + msg.Topic = topicB 137 err = publisher.PublishMessage(msg) 138 require.NoError(t, err) 139 140 cancel() 141 142 select { 143 case <-consumerFinCh: 144 break ··· 147 } 148 149 assert.Len(t, receivedMessages, 1) 150 + assert.Equal(t, topicB, receivedMessages[0].Topic) 151 } 152 153 func TestPublishAndSubscribe(t *testing.T) { ··· 160 sub.Close() 161 }) 162 163 + topics := []string{topicA, topicB} 164 165 err = sub.SubscribeToTopics(topics) 166 require.NoError(t, err) ··· 173 consumer := sub.Consume(ctx) 174 require.NoError(t, err) 175 176 + var receivedMessages []Message 177 178 consumerFinCh := make(chan struct{}) 179 go func() { ··· 185 consumerFinCh <- struct{}{} 186 }() 187 188 + publisher, err := NewPublisher("localhost:9999") 189 require.NoError(t, err) 190 t.Cleanup(func() { 191 publisher.Close() 192 }) 193 194 // send some messages 195 + sentMessages := make([]Message, 0, 10) 196 for i := 0; i < 10; i++ { 197 + msg := Message{ 198 + Topic: topicA, 199 Data: []byte(fmt.Sprintf("message %d", i)), 200 } 201
+28 -65
server/peer.go
··· 1 package server 2 3 import ( 4 - "encoding/binary" 5 - "fmt" 6 "log/slog" 7 "net" 8 - ) 9 - 10 - type peer struct { 11 - conn net.Conn 12 - } 13 - 14 - func newPeer(conn net.Conn) peer { 15 - return peer{ 16 - conn: conn, 17 - } 18 - } 19 20 - // Read wraps the peers underlying connections Read function to satisfy io.Reader 21 - func (p *peer) Read(b []byte) (n int, err error) { 22 - return p.conn.Read(b) 23 - } 24 - 25 - // Write wraps the peers underlying connections Write function to satisfy io.Writer 26 - func (p *peer) Write(b []byte) (n int, err error) { 27 - return p.conn.Write(b) 28 - } 29 - 30 - func (p *peer) addr() net.Addr { 31 - return p.conn.LocalAddr() 32 - } 33 - 34 - func (p *peer) readAction() (Action, error) { 35 - var action Action 36 - err := binary.Read(p.conn, binary.BigEndian, &action) 37 - if err != nil { 38 - return 0, fmt.Errorf("failed to read action from peer: %w", err) 39 - } 40 - 41 - return action, nil 42 - } 43 - 44 - func (p *peer) readDataLength() (uint32, error) { 45 - var dataLen uint32 46 - err := binary.Read(p.conn, binary.BigEndian, &dataLen) 47 - if err != nil { 48 - return 0, fmt.Errorf("failed to read data length from peer: %w", err) 49 - } 50 - 51 - return dataLen, nil 52 - } 53 54 // Status represents the status of a request 55 type Status uint8 ··· 73 return "" 74 } 75 76 - func (p *peer) writeStatus(status Status, message string) { 77 - err := binary.Write(p.conn, binary.BigEndian, status) 78 - if err != nil { 79 - slog.Error("failed to write status to peers connection", "error", err, "peer", p.addr()) 80 - return 81 - } 82 83 - if message == "" { 84 - return 85 } 86 87 - msgBytes := []byte(message) 88 - err = binary.Write(p.conn, binary.BigEndian, uint32(len(msgBytes))) 89 - if err != nil { 90 - slog.Error("failed to write message length to peers connection", "error", err, "peer", p.addr()) 91 - return 92 - } 93 94 - _, err = p.conn.Write(msgBytes) 95 - if err != nil { 96 - slog.Error("failed to write message to peers connection", "error", err, "peer", p.addr()) 97 - return 98 - } 99 }
··· 1 package server 2 3 import ( 4 "log/slog" 5 "net" 6 + "sync" 7 8 + "github.com/google/uuid" 9 + ) 10 11 // Status represents the status of a request 12 type Status uint8 ··· 30 return "" 31 } 32 33 + type peer struct { 34 + conn net.Conn 35 + connMu sync.Mutex 36 + name string 37 + } 38 39 + func newPeer(conn net.Conn) peer { 40 + return peer{ 41 + conn: conn, 42 + name: uuid.New().String(), 43 } 44 + } 45 46 + func (p *peer) addr() net.Addr { 47 + return p.conn.RemoteAddr() 48 + } 49 + 50 + type connOpp func(conn net.Conn) error 51 + 52 + func (p *peer) connOperation(op connOpp, from string) error { 53 + slog.Info("operation running", "from", from, "peer", p.conn.RemoteAddr(), "name", p.name, "mu addr", &p.connMu) 54 + 55 + p.connMu.Lock() 56 + err := op(p.conn) 57 + p.connMu.Unlock() 58 59 + slog.Info("operation finished", "from", from, "peer", p.conn.RemoteAddr(), "name", p.name, "mu addr", &p.connMu) 60 + 61 + return err 62 }
+199 -88
server/server.go
··· 1 package server 2 3 import ( 4 - "context" 5 "encoding/json" 6 "errors" 7 "fmt" 8 "log/slog" 9 "net" 10 "sync" 11 - 12 - "github.com/willdot/messagebroker" 13 ) 14 15 // Action represents the type of action that a peer requests to do ··· 31 } 32 33 // New creates and starts a new server 34 - func New(ctx context.Context, addr string) (*Server, error) { 35 lis, err := net.Listen("tcp", addr) 36 if err != nil { 37 return nil, fmt.Errorf("failed to listen: %w", err) ··· 42 topics: map[string]topic{}, 43 } 44 45 - go srv.start(ctx) 46 47 return srv, nil 48 } ··· 52 return s.lis.Close() 53 } 54 55 - func (s *Server) start(ctx context.Context) { 56 for { 57 conn, err := s.lis.Accept() 58 if err != nil { ··· 70 71 func (s *Server) handleConn(conn net.Conn) { 72 peer := newPeer(conn) 73 - action, err := peer.readAction() 74 if err != nil { 75 slog.Error("failed to read action from peer", "error", err, "peer", peer.addr()) 76 return ··· 85 s.handlePublish(peer) 86 default: 87 slog.Error("unknown action", "action", action, "peer", peer.addr()) 88 - peer.writeStatus(Error, "unknown action") 89 } 90 } 91 92 func (s *Server) handleSubscribe(peer peer) { 93 // subscribe the peer to the topic 94 - s.subscribePeerToTopic(peer) 95 96 // keep handling the peers connection, getting the action from the peer when it wishes to do something else. 97 // once the peers connection ends, it will be unsubscribed from all topics and returned 98 for { 99 - action, err := peer.readAction() 100 if err != nil { 101 // TODO: see if there's a way to check if the peers connection has been ended etc 102 slog.Error("failed to read action from subscriber", "error", err, "peer", peer.addr()) 103 ··· 108 109 switch action { 110 case Subscribe: 111 - s.subscribePeerToTopic(peer) 112 case Unsubscribe: 113 s.handleUnsubscribe(peer) 114 default: 115 slog.Error("unknown action for subscriber", "action", action, "peer", peer.addr()) 116 - peer.writeStatus(Error, "unknown action") 117 continue 118 } 119 } 120 } 121 122 - func (s *Server) subscribePeerToTopic(peer peer) { 123 - // get the topics the peer wishes to subscribe to 124 - dataLen, err := peer.readDataLength() 125 - if err != nil { 126 - slog.Error(err.Error(), "peer", peer.addr()) 127 - peer.writeStatus(Error, "invalid data length of topics provided") 128 - return 129 - } 130 - if dataLen == 0 { 131 - peer.writeStatus(Error, "data length of topics is 0") 132 - return 133 - } 134 - 135 - buf := make([]byte, dataLen) 136 - _, err = peer.Read(buf) 137 - if err != nil { 138 - slog.Error("failed to read subscibers topic data", "error", err, "peer", peer.addr()) 139 - peer.writeStatus(Error, "failed to read topic data") 140 - return 141 - } 142 - 143 - var topics []string 144 - err = json.Unmarshal(buf, &topics) 145 - if err != nil { 146 - slog.Error("failed to unmarshal subscibers topic data", "error", err, "peer", peer.addr()) 147 - peer.writeStatus(Error, "invalid topic data provided") 148 - return 149 - } 150 151 - s.subscribeToTopics(peer, topics) 152 - peer.writeStatus(Subscribed, "") 153 - } 154 155 - func (s *Server) handleUnsubscribe(peer peer) { 156 - // get the topics the peer wishes to unsubscribe from 157 - dataLen, err := peer.readDataLength() 158 - if err != nil { 159 - slog.Error(err.Error(), "peer", peer.addr()) 160 - peer.writeStatus(Error, "invalid data length of topics provided") 161 - return 162 - } 163 - if dataLen == 0 { 164 - peer.writeStatus(Error, "data length of topics is 0") 165 - return 166 - } 167 168 - buf := make([]byte, dataLen) 169 - _, err = peer.Read(buf) 170 - if err != nil { 171 - slog.Error("failed to read subscibers topic data", "error", err, "peer", peer.addr()) 172 - peer.writeStatus(Error, "failed to read topic data") 173 - return 174 - } 175 176 - var topics []string 177 - err = json.Unmarshal(buf, &topics) 178 - if err != nil { 179 - slog.Error("failed to unmarshal subscibers topic data", "error", err, "peer", peer.addr()) 180 - peer.writeStatus(Error, "invalid topic data provided") 181 - return 182 } 183 184 - s.unsubscribeToTopics(peer, topics) 185 - peer.writeStatus(Unsubscribed, "") 186 } 187 188 - func (s *Server) handlePublish(peer peer) { 189 - for { 190 - dataLen, err := peer.readDataLength() 191 if err != nil { 192 slog.Error(err.Error(), "peer", peer.addr()) 193 - peer.writeStatus(Error, "invalid data length of data provided") 194 - return 195 } 196 if dataLen == 0 { 197 - continue 198 } 199 200 buf := make([]byte, dataLen) 201 - _, err = peer.Read(buf) 202 if err != nil { 203 - slog.Error("failed to read data from peer", "error", err, "peer", peer.addr()) 204 - peer.writeStatus(Error, "failed to read data") 205 - return 206 } 207 208 - var msg messagebroker.Message 209 - err = json.Unmarshal(buf, &msg) 210 if err != nil { 211 - slog.Error("failed to unmarshal data to message", "error", err, "peer", peer.addr()) 212 - peer.writeStatus(Error, "invalid message") 213 continue 214 } 215 216 - topic := s.getTopic(msg.Topic) 217 if topic != nil { 218 - topic.sendMessageToSubscribers(msg) 219 } 220 } 221 } 222 223 - func (s *Server) subscribeToTopics(peer peer, topics []string) { 224 for _, topic := range topics { 225 s.addSubsciberToTopic(topic, peer) 226 } 227 } 228 229 - func (s *Server) addSubsciberToTopic(topicName string, peer peer) { 230 s.mu.Lock() 231 defer s.mu.Unlock() 232 ··· 280 281 return nil 282 }
··· 1 package server 2 3 import ( 4 + "encoding/binary" 5 "encoding/json" 6 "errors" 7 "fmt" 8 "log/slog" 9 "net" 10 + "strings" 11 "sync" 12 + "time" 13 ) 14 15 // Action represents the type of action that a peer requests to do ··· 31 } 32 33 // New creates and starts a new server 34 + func New(addr string) (*Server, error) { 35 lis, err := net.Listen("tcp", addr) 36 if err != nil { 37 return nil, fmt.Errorf("failed to listen: %w", err) ··· 42 topics: map[string]topic{}, 43 } 44 45 + go srv.start() 46 47 return srv, nil 48 } ··· 52 return s.lis.Close() 53 } 54 55 + func (s *Server) start() { 56 for { 57 conn, err := s.lis.Accept() 58 if err != nil { ··· 70 71 func (s *Server) handleConn(conn net.Conn) { 72 peer := newPeer(conn) 73 + action, err := readAction(peer) 74 if err != nil { 75 slog.Error("failed to read action from peer", "error", err, "peer", peer.addr()) 76 return ··· 85 s.handlePublish(peer) 86 default: 87 slog.Error("unknown action", "action", action, "peer", peer.addr()) 88 + writeStatus(Error, "unknown action", peer.conn) 89 } 90 } 91 92 func (s *Server) handleSubscribe(peer peer) { 93 // subscribe the peer to the topic 94 + s.subscribePeerToTopic(&peer) 95 96 // keep handling the peers connection, getting the action from the peer when it wishes to do something else. 97 // once the peers connection ends, it will be unsubscribed from all topics and returned 98 for { 99 + action, err := readAction(peer) 100 if err != nil { 101 + var neterr net.Error 102 + if errors.As(err, &neterr) && neterr.Timeout() { 103 + time.Sleep(time.Second) 104 + continue 105 + } 106 // TODO: see if there's a way to check if the peers connection has been ended etc 107 slog.Error("failed to read action from subscriber", "error", err, "peer", peer.addr()) 108 ··· 113 114 switch action { 115 case Subscribe: 116 + s.subscribePeerToTopic(&peer) 117 case Unsubscribe: 118 s.handleUnsubscribe(peer) 119 default: 120 slog.Error("unknown action for subscriber", "action", action, "peer", peer.addr()) 121 + writeStatus(Error, "unknown action", peer.conn) 122 continue 123 } 124 } 125 } 126 127 + func (s *Server) subscribePeerToTopic(peer *peer) { 128 + op := func(conn net.Conn) error { 129 + // get the topics the peer wishes to subscribe to 130 + dataLen, err := dataLength(conn) 131 + if err != nil { 132 + slog.Error(err.Error(), "peer", peer.addr()) 133 + writeStatus(Error, "invalid data length of topics provided", conn) 134 + return nil 135 + } 136 + if dataLen == 0 { 137 + writeStatus(Error, "data length of topics is 0", conn) 138 + return nil 139 + } 140 141 + buf := make([]byte, dataLen) 142 + _, err = conn.Read(buf) 143 + if err != nil { 144 + slog.Error("failed to read subscibers topic data", "error", err, "peer", peer.addr()) 145 + writeStatus(Error, "failed to read topic data", conn) 146 + return nil 147 + } 148 149 + var topics []string 150 + err = json.Unmarshal(buf, &topics) 151 + if err != nil { 152 + slog.Error("failed to unmarshal subscibers topic data", "error", err, "peer", peer.addr()) 153 + writeStatus(Error, "invalid topic data provided", conn) 154 + return nil 155 + } 156 157 + s.subscribeToTopics(peer, topics) 158 + writeStatus(Subscribed, "", conn) 159 160 + return nil 161 } 162 163 + _ = peer.connOperation(op, "subscribe peer to topic") 164 } 165 166 + func (s *Server) handleUnsubscribe(peer peer) { 167 + op := func(conn net.Conn) error { 168 + // get the topics the peer wishes to unsubscribe from 169 + dataLen, err := dataLength(conn) 170 if err != nil { 171 slog.Error(err.Error(), "peer", peer.addr()) 172 + writeStatus(Error, "invalid data length of topics provided", conn) 173 + return nil 174 } 175 if dataLen == 0 { 176 + writeStatus(Error, "data length of topics is 0", conn) 177 + return nil 178 } 179 180 buf := make([]byte, dataLen) 181 + _, err = conn.Read(buf) 182 if err != nil { 183 + slog.Error("failed to read subscibers topic data", "error", err, "peer", peer.addr()) 184 + writeStatus(Error, "failed to read topic data", conn) 185 + return nil 186 } 187 188 + var topics []string 189 + err = json.Unmarshal(buf, &topics) 190 if err != nil { 191 + slog.Error("failed to unmarshal subscibers topic data", "error", err, "peer", peer.addr()) 192 + writeStatus(Error, "invalid topic data provided", conn) 193 + return nil 194 + } 195 + 196 + s.unsubscribeToTopics(peer, topics) 197 + writeStatus(Unsubscribed, "", conn) 198 + 199 + return nil 200 + } 201 + 202 + _ = peer.connOperation(op, "handle unsubscribe") 203 + } 204 + 205 + type messageToSend struct { 206 + topic string 207 + data []byte 208 + } 209 + 210 + func (s *Server) handlePublish(peer peer) { 211 + for { 212 + var message *messageToSend 213 + 214 + op := func(conn net.Conn) error { 215 + dataLen, err := dataLength(conn) 216 + if err != nil { 217 + slog.Error(err.Error(), "peer", peer.addr()) 218 + writeStatus(Error, "invalid data length of data provided", conn) 219 + return nil 220 + } 221 + if dataLen == 0 { 222 + return nil 223 + } 224 + topicBuf := make([]byte, dataLen) 225 + _, err = conn.Read(topicBuf) 226 + if err != nil { 227 + slog.Error("failed to read topic from peer", "error", err, "peer", peer.addr()) 228 + writeStatus(Error, "failed to read topic", conn) 229 + return nil 230 + } 231 + 232 + topicStr := string(topicBuf) 233 + if !strings.HasPrefix(topicStr, "topic:") { 234 + slog.Error("topic data does not contain topic prefix", "peer", peer.addr()) 235 + writeStatus(Error, "topic data does not contain 'topic:' prefix", conn) 236 + return nil 237 + } 238 + topicStr = strings.TrimPrefix(topicStr, "topic:") 239 + 240 + dataLen, err = dataLength(conn) 241 + if err != nil { 242 + slog.Error(err.Error(), "peer", peer.addr()) 243 + writeStatus(Error, "invalid data length of data provided", conn) 244 + return nil 245 + } 246 + if dataLen == 0 { 247 + return nil 248 + } 249 + 250 + dataBuf := make([]byte, dataLen) 251 + _, err = conn.Read(dataBuf) 252 + if err != nil { 253 + slog.Error("failed to read data from peer", "error", err, "peer", peer.addr()) 254 + writeStatus(Error, "failed to read data", conn) 255 + return nil 256 + } 257 + 258 + message = &messageToSend{ 259 + topic: topicStr, 260 + data: dataBuf, 261 + } 262 + return nil 263 + } 264 + 265 + _ = peer.connOperation(op, "handle publish") 266 + 267 + if message == nil { 268 continue 269 } 270 + // TODO: this can be done in a go routine because once we've got the message from the publisher, the publisher 271 + // doesn't need to wait for us to send the message to all peers 272 273 + topic := s.getTopic(message.topic) 274 if topic != nil { 275 + topic.sendMessageToSubscribers(message.data) 276 } 277 } 278 } 279 280 + func (s *Server) subscribeToTopics(peer *peer, topics []string) { 281 for _, topic := range topics { 282 s.addSubsciberToTopic(topic, peer) 283 } 284 } 285 286 + func (s *Server) addSubsciberToTopic(topicName string, peer *peer) { 287 s.mu.Lock() 288 defer s.mu.Unlock() 289 ··· 337 338 return nil 339 } 340 + 341 + func readAction(peer peer) (Action, error) { 342 + var action Action 343 + op := func(conn net.Conn) error { 344 + conn.SetReadDeadline(time.Now().Add(time.Second)) 345 + 346 + err := binary.Read(conn, binary.BigEndian, &action) 347 + if err != nil { 348 + return err 349 + } 350 + return nil 351 + } 352 + 353 + err := peer.connOperation(op, "read action") 354 + if err != nil { 355 + return 0, fmt.Errorf("failed to read action from peer: %w", err) 356 + } 357 + 358 + return action, nil 359 + } 360 + 361 + func dataLength(conn net.Conn) (uint32, error) { 362 + var dataLen uint32 363 + err := binary.Read(conn, binary.BigEndian, &dataLen) 364 + if err != nil { 365 + return 0, err 366 + } 367 + return dataLen, nil 368 + } 369 + 370 + func writeStatus(status Status, message string, conn net.Conn) { 371 + err := binary.Write(conn, binary.BigEndian, status) 372 + if err != nil { 373 + slog.Error("failed to write status to peers connection", "error", err, "peer", conn.RemoteAddr()) 374 + return 375 + } 376 + 377 + if message == "" { 378 + return 379 + } 380 + 381 + msgBytes := []byte(message) 382 + err = binary.Write(conn, binary.BigEndian, uint32(len(msgBytes))) 383 + if err != nil { 384 + slog.Error("failed to write message length to peers connection", "error", err, "peer", conn.RemoteAddr()) 385 + return 386 + } 387 + 388 + _, err = conn.Write(msgBytes) 389 + if err != nil { 390 + slog.Error("failed to write message to peers connection", "error", err, "peer", conn.RemoteAddr()) 391 + return 392 + } 393 + }
+82 -62
server/server_test.go
··· 1 package server 2 3 import ( 4 - "context" 5 "encoding/binary" 6 "encoding/json" 7 "fmt" ··· 11 12 "github.com/stretchr/testify/assert" 13 "github.com/stretchr/testify/require" 14 - "github.com/willdot/messagebroker" 15 ) 16 17 func createServer(t *testing.T) *Server { 18 - srv, err := New(context.Background(), ":3000") 19 require.NoError(t, err) 20 21 t.Cleanup(func() { ··· 36 } 37 38 func createConnectionAndSubscribe(t *testing.T, topics []string) net.Conn { 39 - conn, err := net.Dial("tcp", "localhost:3000") 40 require.NoError(t, err) 41 42 err = binary.Write(conn, binary.BigEndian, Subscribe) ··· 64 func TestSubscribeToTopics(t *testing.T) { 65 // create a server with an existing topic so we can test subscribing to a new and 66 // existing topic 67 - srv := createServerWithExistingTopic(t, "topic a") 68 69 - _ = createConnectionAndSubscribe(t, []string{"topic a", "topic b"}) 70 71 assert.Len(t, srv.topics, 2) 72 - assert.Len(t, srv.topics["topic a"].subscriptions, 1) 73 - assert.Len(t, srv.topics["topic b"].subscriptions, 1) 74 } 75 76 func TestUnsubscribesFromTopic(t *testing.T) { 77 - srv := createServerWithExistingTopic(t, "topic a") 78 79 - conn := createConnectionAndSubscribe(t, []string{"topic a", "topic b", "topic c"}) 80 81 assert.Len(t, srv.topics, 3) 82 - assert.Len(t, srv.topics["topic a"].subscriptions, 1) 83 - assert.Len(t, srv.topics["topic b"].subscriptions, 1) 84 - assert.Len(t, srv.topics["topic c"].subscriptions, 1) 85 86 err := binary.Write(conn, binary.BigEndian, Unsubscribe) 87 require.NoError(t, err) 88 89 - topics := []string{"topic a", "topic b"} 90 rawTopics, err := json.Marshal(topics) 91 require.NoError(t, err) 92 ··· 104 assert.Equal(t, expectedRes, int(resp)) 105 106 assert.Len(t, srv.topics, 3) 107 - assert.Len(t, srv.topics["topic a"].subscriptions, 0) 108 - assert.Len(t, srv.topics["topic b"].subscriptions, 0) 109 - assert.Len(t, srv.topics["topic c"].subscriptions, 1) 110 } 111 112 func TestSubscriberClosesWithoutUnsubscribing(t *testing.T) { 113 srv := createServer(t) 114 115 - conn := createConnectionAndSubscribe(t, []string{"topic a", "topic b"}) 116 117 assert.Len(t, srv.topics, 2) 118 - assert.Len(t, srv.topics["topic a"].subscriptions, 1) 119 - assert.Len(t, srv.topics["topic b"].subscriptions, 1) 120 121 // close the conn 122 err := conn.Close() 123 require.NoError(t, err) 124 125 - publisherConn, err := net.Dial("tcp", "localhost:3000") 126 require.NoError(t, err) 127 128 err = binary.Write(publisherConn, binary.BigEndian, Publish) ··· 137 require.Equal(t, len(data), n) 138 139 assert.Len(t, srv.topics, 2) 140 - assert.Len(t, srv.topics["topic a"].subscriptions, 0) 141 - assert.Len(t, srv.topics["topic b"].subscriptions, 0) 142 } 143 144 func TestInvalidAction(t *testing.T) { 145 _ = createServer(t) 146 147 - conn, err := net.Dial("tcp", "localhost:3000") 148 require.NoError(t, err) 149 150 err = binary.Write(conn, binary.BigEndian, uint8(99)) ··· 170 assert.Equal(t, expectedMessage, string(buf)) 171 } 172 173 - func TestInvalidMessagePublished(t *testing.T) { 174 _ = createServer(t) 175 176 - publisherConn, err := net.Dial("tcp", "localhost:3000") 177 require.NoError(t, err) 178 179 err = binary.Write(publisherConn, binary.BigEndian, Publish) 180 require.NoError(t, err) 181 182 - // send some data 183 - data := []byte("this isn't wrapped in a message type") 184 - 185 - // send data length first 186 - err = binary.Write(publisherConn, binary.BigEndian, uint32(len(data))) 187 require.NoError(t, err) 188 - n, err := publisherConn.Write(data) 189 require.NoError(t, err) 190 - require.Equal(t, len(data), n) 191 192 expectedRes := Error 193 ··· 196 197 assert.Equal(t, expectedRes, int(resp)) 198 199 - expectedMessage := "invalid message" 200 201 var dataLen uint32 202 err = binary.Read(publisherConn, binary.BigEndian, &dataLen) ··· 212 func TestSendsDataToTopicSubscribers(t *testing.T) { 213 _ = createServer(t) 214 215 - subscribers := make([]net.Conn, 0, 5) 216 - for i := 0; i < 5; i++ { 217 - subscriberConn := createConnectionAndSubscribe(t, []string{"topic a", "topic b"}) 218 219 subscribers = append(subscribers, subscriberConn) 220 } 221 222 - publisherConn, err := net.Dial("tcp", "localhost:3000") 223 require.NoError(t, err) 224 225 err = binary.Write(publisherConn, binary.BigEndian, Publish) 226 require.NoError(t, err) 227 228 - // send a message 229 - msg := messagebroker.Message{ 230 - Topic: "topic a", 231 - Data: []byte("hello world"), 232 - } 233 234 - rawMsg, err := json.Marshal(msg) 235 require.NoError(t, err) 236 237 - // send data length first 238 - err = binary.Write(publisherConn, binary.BigEndian, uint32(len(rawMsg))) 239 require.NoError(t, err) 240 - n, err := publisherConn.Write(rawMsg) 241 require.NoError(t, err) 242 - require.Equal(t, len(rawMsg), n) 243 244 // check the subsribers got the data 245 for _, conn := range subscribers { 246 247 var dataLen uint64 248 err = binary.Read(conn, binary.BigEndian, &dataLen) ··· 253 require.NoError(t, err) 254 require.Equal(t, int(dataLen), n) 255 256 - assert.Equal(t, rawMsg, buf) 257 } 258 } 259 260 func TestPublishMultipleTimes(t *testing.T) { 261 _ = createServer(t) 262 263 - publisherConn, err := net.Dial("tcp", "localhost:3000") 264 require.NoError(t, err) 265 266 err = binary.Write(publisherConn, binary.BigEndian, Publish) ··· 268 269 messages := make([][]byte, 0, 10) 270 for i := 0; i < 10; i++ { 271 - msg := messagebroker.Message{ 272 - Topic: "topic a", 273 - Data: []byte(fmt.Sprintf("message %d", i)), 274 - } 275 - 276 - rawMsg, err := json.Marshal(msg) 277 - require.NoError(t, err) 278 - 279 - messages = append(messages, rawMsg) 280 } 281 282 subscribeFinCh := make(chan struct{}) 283 // create a subscriber that will read messages 284 - subscriberConn := createConnectionAndSubscribe(t, []string{"topic a", "topic b"}) 285 go func() { 286 // check subscriber got all messages 287 for _, msg := range messages { 288 var dataLen uint64 289 err = binary.Read(subscriberConn, binary.BigEndian, &dataLen) 290 require.NoError(t, err) ··· 300 subscribeFinCh <- struct{}{} 301 }() 302 303 // send multiple messages 304 for _, msg := range messages { 305 - // send data length first 306 err = binary.Write(publisherConn, binary.BigEndian, uint32(len(msg))) 307 require.NoError(t, err) 308 - n, err := publisherConn.Write(msg) 309 require.NoError(t, err) 310 require.Equal(t, len(msg), n) 311 }
··· 1 package server 2 3 import ( 4 "encoding/binary" 5 "encoding/json" 6 "fmt" ··· 10 11 "github.com/stretchr/testify/assert" 12 "github.com/stretchr/testify/require" 13 + ) 14 + 15 + const ( 16 + topicA = "topic a" 17 + topicB = "topic b" 18 + topicC = "topic c" 19 + 20 + serverAddr = ":6666" 21 ) 22 23 func createServer(t *testing.T) *Server { 24 + srv, err := New(serverAddr) 25 require.NoError(t, err) 26 27 t.Cleanup(func() { ··· 42 } 43 44 func createConnectionAndSubscribe(t *testing.T, topics []string) net.Conn { 45 + conn, err := net.Dial("tcp", fmt.Sprintf("localhost%s", serverAddr)) 46 require.NoError(t, err) 47 48 err = binary.Write(conn, binary.BigEndian, Subscribe) ··· 70 func TestSubscribeToTopics(t *testing.T) { 71 // create a server with an existing topic so we can test subscribing to a new and 72 // existing topic 73 + srv := createServerWithExistingTopic(t, topicA) 74 75 + _ = createConnectionAndSubscribe(t, []string{topicA, topicB}) 76 77 assert.Len(t, srv.topics, 2) 78 + assert.Len(t, srv.topics[topicA].subscriptions, 1) 79 + assert.Len(t, srv.topics[topicB].subscriptions, 1) 80 } 81 82 func TestUnsubscribesFromTopic(t *testing.T) { 83 + srv := createServerWithExistingTopic(t, topicA) 84 85 + conn := createConnectionAndSubscribe(t, []string{topicA, topicB, topicC}) 86 87 assert.Len(t, srv.topics, 3) 88 + assert.Len(t, srv.topics[topicA].subscriptions, 1) 89 + assert.Len(t, srv.topics[topicB].subscriptions, 1) 90 + assert.Len(t, srv.topics[topicC].subscriptions, 1) 91 92 err := binary.Write(conn, binary.BigEndian, Unsubscribe) 93 require.NoError(t, err) 94 95 + topics := []string{topicA, topicB} 96 rawTopics, err := json.Marshal(topics) 97 require.NoError(t, err) 98 ··· 110 assert.Equal(t, expectedRes, int(resp)) 111 112 assert.Len(t, srv.topics, 3) 113 + assert.Len(t, srv.topics[topicA].subscriptions, 0) 114 + assert.Len(t, srv.topics[topicB].subscriptions, 0) 115 + assert.Len(t, srv.topics[topicC].subscriptions, 1) 116 } 117 118 func TestSubscriberClosesWithoutUnsubscribing(t *testing.T) { 119 srv := createServer(t) 120 121 + conn := createConnectionAndSubscribe(t, []string{topicA, topicB}) 122 123 assert.Len(t, srv.topics, 2) 124 + assert.Len(t, srv.topics[topicA].subscriptions, 1) 125 + assert.Len(t, srv.topics[topicB].subscriptions, 1) 126 127 // close the conn 128 err := conn.Close() 129 require.NoError(t, err) 130 131 + publisherConn, err := net.Dial("tcp", fmt.Sprintf("localhost%s", serverAddr)) 132 require.NoError(t, err) 133 134 err = binary.Write(publisherConn, binary.BigEndian, Publish) ··· 143 require.Equal(t, len(data), n) 144 145 assert.Len(t, srv.topics, 2) 146 + assert.Len(t, srv.topics[topicA].subscriptions, 0) 147 + assert.Len(t, srv.topics[topicB].subscriptions, 0) 148 } 149 150 func TestInvalidAction(t *testing.T) { 151 _ = createServer(t) 152 153 + conn, err := net.Dial("tcp", fmt.Sprintf("localhost%s", serverAddr)) 154 require.NoError(t, err) 155 156 err = binary.Write(conn, binary.BigEndian, uint8(99)) ··· 176 assert.Equal(t, expectedMessage, string(buf)) 177 } 178 179 + func TestInvalidTopicDataPublished(t *testing.T) { 180 _ = createServer(t) 181 182 + publisherConn, err := net.Dial("tcp", fmt.Sprintf("localhost%s", serverAddr)) 183 require.NoError(t, err) 184 185 err = binary.Write(publisherConn, binary.BigEndian, Publish) 186 require.NoError(t, err) 187 188 + // send topic 189 + topic := topicA 190 + err = binary.Write(publisherConn, binary.BigEndian, uint32(len(topic))) 191 require.NoError(t, err) 192 + _, err = publisherConn.Write([]byte(topic)) 193 require.NoError(t, err) 194 195 expectedRes := Error 196 ··· 199 200 assert.Equal(t, expectedRes, int(resp)) 201 202 + expectedMessage := "topic data does not contain 'topic:' prefix" 203 204 var dataLen uint32 205 err = binary.Read(publisherConn, binary.BigEndian, &dataLen) ··· 215 func TestSendsDataToTopicSubscribers(t *testing.T) { 216 _ = createServer(t) 217 218 + subscribers := make([]net.Conn, 0, 1) 219 + for i := 0; i < 1; i++ { 220 + subscriberConn := createConnectionAndSubscribe(t, []string{topicA, topicB}) 221 222 subscribers = append(subscribers, subscriberConn) 223 } 224 225 + publisherConn, err := net.Dial("tcp", fmt.Sprintf("localhost%s", serverAddr)) 226 require.NoError(t, err) 227 228 err = binary.Write(publisherConn, binary.BigEndian, Publish) 229 require.NoError(t, err) 230 231 + topic := fmt.Sprintf("topic:%s", topicA) 232 + messageData := "hello world" 233 234 + // send topic first 235 + err = binary.Write(publisherConn, binary.BigEndian, uint32(len(topic))) 236 + require.NoError(t, err) 237 + _, err = publisherConn.Write([]byte(topic)) 238 require.NoError(t, err) 239 240 + // now send the data 241 + err = binary.Write(publisherConn, binary.BigEndian, uint32(len(messageData))) 242 require.NoError(t, err) 243 + n, err := publisherConn.Write([]byte(messageData)) 244 require.NoError(t, err) 245 + require.Equal(t, len(messageData), n) 246 247 // check the subsribers got the data 248 for _, conn := range subscribers { 249 + var topicLen uint64 250 + err = binary.Read(conn, binary.BigEndian, &topicLen) 251 + require.NoError(t, err) 252 + 253 + topicBuf := make([]byte, topicLen) 254 + _, err = conn.Read(topicBuf) 255 + require.NoError(t, err) 256 + assert.Equal(t, topicA, string(topicBuf)) 257 258 var dataLen uint64 259 err = binary.Read(conn, binary.BigEndian, &dataLen) ··· 264 require.NoError(t, err) 265 require.Equal(t, int(dataLen), n) 266 267 + assert.Equal(t, messageData, string(buf)) 268 } 269 } 270 271 func TestPublishMultipleTimes(t *testing.T) { 272 _ = createServer(t) 273 274 + publisherConn, err := net.Dial("tcp", fmt.Sprintf("localhost%s", serverAddr)) 275 require.NoError(t, err) 276 277 err = binary.Write(publisherConn, binary.BigEndian, Publish) ··· 279 280 messages := make([][]byte, 0, 10) 281 for i := 0; i < 10; i++ { 282 + messages = append(messages, []byte(fmt.Sprintf("message %d", i))) 283 } 284 285 subscribeFinCh := make(chan struct{}) 286 // create a subscriber that will read messages 287 + subscriberConn := createConnectionAndSubscribe(t, []string{topicA, topicB}) 288 go func() { 289 // check subscriber got all messages 290 for _, msg := range messages { 291 + var topicLen uint64 292 + err = binary.Read(subscriberConn, binary.BigEndian, &topicLen) 293 + require.NoError(t, err) 294 + 295 + topicBuf := make([]byte, topicLen) 296 + _, err = subscriberConn.Read(topicBuf) 297 + require.NoError(t, err) 298 + assert.Equal(t, topicA, string(topicBuf)) 299 + 300 var dataLen uint64 301 err = binary.Read(subscriberConn, binary.BigEndian, &dataLen) 302 require.NoError(t, err) ··· 312 subscribeFinCh <- struct{}{} 313 }() 314 315 + topic := fmt.Sprintf("topic:%s", topicA) 316 + 317 // send multiple messages 318 for _, msg := range messages { 319 + // send topic first 320 + err = binary.Write(publisherConn, binary.BigEndian, uint32(len(topic))) 321 + require.NoError(t, err) 322 + _, err = publisherConn.Write([]byte(topic)) 323 + require.NoError(t, err) 324 + 325 + // now send the data 326 err = binary.Write(publisherConn, binary.BigEndian, uint32(len(msg))) 327 require.NoError(t, err) 328 + n, err := publisherConn.Write([]byte(msg)) 329 require.NoError(t, err) 330 require.Equal(t, len(msg), n) 331 }
-26
server/subscriber.go
··· 1 - package server 2 - 3 - import ( 4 - "encoding/binary" 5 - "fmt" 6 - ) 7 - 8 - type subscriber struct { 9 - peer peer 10 - currentOffset int 11 - } 12 - 13 - func (s *subscriber) sendMessage(msg []byte) error { 14 - dataLen := uint64(len(msg)) 15 - 16 - err := binary.Write(&s.peer, binary.BigEndian, dataLen) 17 - if err != nil { 18 - return fmt.Errorf("failed to send data length: %w", err) 19 - } 20 - 21 - _, err = s.peer.Write(msg) 22 - if err != nil { 23 - return fmt.Errorf("failed to write to peer: %w", err) 24 - } 25 - return nil 26 - }
···
+39 -11
server/topic.go
··· 1 package server 2 3 import ( 4 - "encoding/json" 5 "log/slog" 6 "net" 7 "sync" 8 - 9 - "github.com/willdot/messagebroker" 10 ) 11 12 type topic struct { ··· 15 mu sync.Mutex 16 } 17 18 func newTopic(name string) topic { 19 return topic{ 20 name: name, ··· 30 delete(t.subscriptions, addr) 31 } 32 33 - func (t *topic) sendMessageToSubscribers(msg messagebroker.Message) { 34 t.mu.Lock() 35 subscribers := t.subscriptions 36 t.mu.Unlock() 37 38 - msgData, err := json.Marshal(msg) 39 - if err != nil { 40 - slog.Error("failed to marshal message for subscribers", "error", err) 41 - } 42 43 - for addr, subscriber := range subscribers { 44 - err := subscriber.sendMessage(msgData) 45 if err != nil { 46 slog.Error("failed to send to message", "error", err, "peer", addr) 47 - continue 48 } 49 } 50 }
··· 1 package server 2 3 import ( 4 + "encoding/binary" 5 + "fmt" 6 "log/slog" 7 "net" 8 "sync" 9 ) 10 11 type topic struct { ··· 14 mu sync.Mutex 15 } 16 17 + type subscriber struct { 18 + peer *peer 19 + currentOffset int 20 + } 21 + 22 func newTopic(name string) topic { 23 return topic{ 24 name: name, ··· 34 delete(t.subscriptions, addr) 35 } 36 37 + func (t *topic) sendMessageToSubscribers(msgData []byte) { 38 t.mu.Lock() 39 subscribers := t.subscriptions 40 t.mu.Unlock() 41 42 + for addr, subscriber := range subscribers { 43 + //sendMessageOpFunc := sendMessageOp(t.name, msgData) 44 45 + err := subscriber.peer.connOperation(sendMessageOp(t.name, msgData), "send message to subscribers") 46 if err != nil { 47 slog.Error("failed to send to message", "error", err, "peer", addr) 48 + return 49 } 50 } 51 } 52 + 53 + func sendMessageOp(topic string, data []byte) connOpp { 54 + return func(conn net.Conn) error { 55 + topicLen := uint64(len(topic)) 56 + err := binary.Write(conn, binary.BigEndian, topicLen) 57 + if err != nil { 58 + return fmt.Errorf("failed to send topic length: %w", err) 59 + } 60 + _, err = conn.Write([]byte(topic)) 61 + if err != nil { 62 + return fmt.Errorf("failed to send topic: %w", err) 63 + } 64 + 65 + dataLen := uint64(len(data)) 66 + 67 + err = binary.Write(conn, binary.BigEndian, dataLen) 68 + if err != nil { 69 + return fmt.Errorf("failed to send data length: %w", err) 70 + } 71 + 72 + _, err = conn.Write(data) 73 + if err != nil { 74 + return fmt.Errorf("failed to write to peer: %w", err) 75 + } 76 + return nil 77 + } 78 + }