+20
dockerfile.example-server
+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
+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
+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
+5
-1
go.mod
+4
go.sum
+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
-1
message.go
pubsub/message.go
+33
-16
pubsub/publisher.go
+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
+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
+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
+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
+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
+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
-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
+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
+
}