1package pubsub
  2
  3import (
  4	"errors"
  5	"io"
  6	"iter"
  7	"log/slog"
  8	"path"
  9	"reflect"
 10	"strings"
 11	"sync"
 12	"time"
 13
 14	"github.com/antoniomika/syncmap"
 15)
 16
 17// HasWildcard checks if a topic string contains the wildcard character (*).
 18func HasWildcard(topic string) bool {
 19	return strings.Contains(topic, "*")
 20}
 21
 22// MatchTopic returns true if pattern matches topic exactly or via path.Match wildcarding.
 23func MatchTopic(pattern, topic string) bool {
 24	if pattern == topic {
 25		return true
 26	}
 27	if HasWildcard(pattern) {
 28		matched, err := path.Match(pattern, topic)
 29		return err == nil && matched
 30	}
 31	return false
 32}
 33
 34/*
 35Broker receives published messages and dispatches the message to the
 36subscribing clients. An message contains a message topic that clients
 37subscribe to and brokers use these subscription lists for determining the
 38clients to receive the message.
 39*/
 40type Broker interface {
 41	GetChannels() iter.Seq2[string, *Channel]
 42	GetClients() iter.Seq2[string, *Client]
 43	Connect(*Client, []*Channel) (error, error)
 44	SetDispatcher(dispatcher MessageDispatcher, channels []*Channel) error
 45}
 46
 47type BaseBroker struct {
 48	Channels *syncmap.Map[string, *Channel]
 49	Logger   *slog.Logger
 50}
 51
 52func (b *BaseBroker) Cleanup() {
 53	toRemove := []string{}
 54	for _, channel := range b.GetChannels() {
 55		count := 0
 56
 57		for range channel.GetClients() {
 58			count++
 59		}
 60
 61		if count == 0 {
 62			channel.Cleanup()
 63			toRemove = append(toRemove, channel.Topic)
 64		}
 65	}
 66
 67	for _, channel := range toRemove {
 68		b.Channels.Delete(channel)
 69	}
 70}
 71
 72func (b *BaseBroker) GetChannels() iter.Seq2[string, *Channel] {
 73	return b.Channels.Range
 74}
 75
 76func (b *BaseBroker) GetClients() iter.Seq2[string, *Client] {
 77	return func(yield func(string, *Client) bool) {
 78		for _, channel := range b.GetChannels() {
 79			channel.Clients.Range(yield)
 80		}
 81	}
 82}
 83
 84func (b *BaseBroker) Connect(client *Client, channels []*Channel) (error, error) {
 85	for _, channel := range channels {
 86		dataChannel := b.ensureChannel(channel)
 87		dataChannel.Clients.Store(client.ID, client)
 88		client.Channels.Store(dataChannel.Topic, dataChannel)
 89
 90		// If client is a subscriber and channel.Topic is a wildcard pattern,
 91		// attach client to all existing concrete channels matching the pattern.
 92		if (client.Direction == ChannelDirectionOutput || client.Direction == ChannelDirectionInputOutput) && HasWildcard(channel.Topic) {
 93			for _, existingChannel := range b.GetChannels() {
 94				if existingChannel.Topic != channel.Topic && !HasWildcard(existingChannel.Topic) && MatchTopic(channel.Topic, existingChannel.Topic) {
 95					existingChannel.Clients.Store(client.ID, client)
 96					client.Channels.Store(existingChannel.Topic, existingChannel)
 97				}
 98			}
 99		}
100
101		defer func() {
102			for _, ch := range client.GetChannels() {
103				ch.Clients.Delete(client.ID)
104				client.Channels.Delete(ch.Topic)
105			}
106			client.Channels.Delete(channel.Topic)
107			dataChannel.Clients.Delete(client.ID)
108
109			client.Cleanup()
110
111			inputCount := 0
112			pipeCount := 0
113			for _, cl := range dataChannel.GetClients() {
114				switch cl.Direction {
115				case ChannelDirectionInput:
116					inputCount++
117				case ChannelDirectionInputOutput:
118					pipeCount++
119				}
120			}
121
122			if inputCount == 0 && pipeCount <= 1 {
123				for _, cl := range dataChannel.GetClients() {
124					if !cl.KeepAlive {
125						otherChannels := 0
126						for _, ch := range cl.GetChannels() {
127							if ch.Topic != dataChannel.Topic {
128								otherChannels++
129							}
130						}
131						if otherChannels == 0 {
132							cl.Cleanup()
133						}
134					}
135				}
136			}
137
138			b.Cleanup()
139		}()
140	}
141
142	var (
143		inputErr  error
144		outputErr error
145		wg        sync.WaitGroup
146	)
147
148	// Pub
149	if client.Direction == ChannelDirectionInput || client.Direction == ChannelDirectionInputOutput {
150		wg.Add(1)
151		go func() {
152			defer wg.Done()
153			for {
154				data := make([]byte, 32*1024)
155				n, err := client.ReadWriter.Read(data)
156
157				data = data[:n]
158
159				channelMessage := ChannelMessage{
160					Data:      data,
161					ClientID:  client.ID,
162					Direction: ChannelDirectionInput,
163				}
164
165				if client.BlockWrite {
166				mainLoop:
167					for {
168						count := 0
169						for _, channel := range client.GetChannels() {
170							for _, chanClient := range channel.GetClients() {
171								if chanClient.Direction == ChannelDirectionOutput || chanClient.Direction == ChannelDirectionInputOutput {
172									count++
173								}
174							}
175						}
176
177						if count > 0 {
178							break mainLoop
179						}
180
181						select {
182						case <-client.Done:
183							break mainLoop
184						case <-time.After(1 * time.Millisecond):
185							continue
186						}
187					}
188				}
189
190				var sendwg sync.WaitGroup
191
192				for _, channel := range client.GetChannels() {
193					sendwg.Add(1)
194					go func() {
195						defer sendwg.Done()
196						select {
197						case channel.Data <- channelMessage:
198						case <-client.Done:
199						case <-channel.Done:
200						}
201					}()
202				}
203
204				sendwg.Wait()
205
206				if err != nil {
207					client.Cleanup()
208					if errors.Is(err, io.EOF) {
209						return
210					}
211					inputErr = err
212					return
213				}
214			}
215		}()
216	}
217
218	// Sub
219	if client.Direction == ChannelDirectionOutput || client.Direction == ChannelDirectionInputOutput {
220		wg.Add(1)
221		go func() {
222			defer wg.Done()
223		mainLoop:
224			for {
225				select {
226				case data, ok := <-client.Data:
227					_, err := client.ReadWriter.Write(data.Data)
228					if err != nil {
229						outputErr = err
230						client.Cleanup()
231						break mainLoop
232					}
233
234					if !ok {
235						break mainLoop
236					}
237				case <-client.Done:
238					break mainLoop
239				}
240			}
241		}()
242	}
243
244	done := make(chan struct{})
245	go func() {
246		wg.Wait()
247		close(done)
248	}()
249
250	select {
251	case <-done:
252	case <-client.Done:
253	}
254
255	return inputErr, outputErr
256}
257
258func (b *BaseBroker) ensureChannel(channel *Channel) *Channel {
259	dataChannel, loaded := b.Channels.LoadOrStore(channel.Topic, channel)
260	dataChannel.Handle()
261
262	// If this is a concrete (non-wildcard) channel created for the first time,
263	// attach any active wildcard subscribers whose pattern matches dataChannel.Topic.
264	if !loaded && !HasWildcard(channel.Topic) {
265		for _, existingChannel := range b.GetChannels() {
266			if HasWildcard(existingChannel.Topic) && MatchTopic(existingChannel.Topic, channel.Topic) {
267				for _, client := range existingChannel.GetClients() {
268					if client.Direction == ChannelDirectionOutput || client.Direction == ChannelDirectionInputOutput {
269						dataChannel.Clients.Store(client.ID, client)
270						client.Channels.Store(dataChannel.Topic, dataChannel)
271					}
272				}
273			}
274		}
275	}
276
277	return dataChannel
278}
279
280func (b *BaseBroker) SetDispatcher(dispatcher MessageDispatcher, channels []*Channel) error {
281	for _, channel := range channels {
282		dataChannel := b.ensureChannel(channel)
283		existingDispatcher := dataChannel.GetDispatcher()
284		if reflect.TypeOf(existingDispatcher) != reflect.TypeOf(dispatcher) {
285			dataChannel.SetDispatcher(dispatcher)
286		}
287	}
288	return nil
289}
290
291var _ Broker = (*BaseBroker)(nil)