├── .dockerignore ├── .github ├── docker-compose.gha.yml └── workflows │ ├── build.yml │ └── test.yml ├── .gitignore ├── .vscode └── launch.json ├── Dockerfile ├── LICENSE ├── Makefile ├── README.md ├── broker.go ├── channel.go ├── client.go ├── cmd ├── basic │ └── main.go └── example │ ├── go.mod │ ├── go.sum │ └── main.go ├── docker-compose.yml ├── go.mod ├── go.sum ├── multicast.go ├── multicast_test.go └── pubsub.go /.dockerignore: -------------------------------------------------------------------------------- 1 | **/ssh_data/ 2 | .git/ 3 | .github/ 4 | .vscode/ 5 | .coverage.out -------------------------------------------------------------------------------- /.github/docker-compose.gha.yml: -------------------------------------------------------------------------------- 1 | services: 2 | pubsub: 3 | build: 4 | cache_from: 5 | - type=gha 6 | cache_to: 7 | - type=gha,mode=max 8 | -------------------------------------------------------------------------------- /.github/workflows/build.yml: -------------------------------------------------------------------------------- 1 | name: Build 2 | 3 | on: 4 | push: 5 | branches: 6 | - main 7 | 8 | jobs: 9 | build: 10 | runs-on: ubuntu-latest 11 | steps: 12 | - uses: actions/checkout@v4 13 | with: 14 | fetch-depth: 0 15 | - name: Expose GitHub Runtime 16 | uses: crazy-max/ghaction-github-runtime@v3 17 | - name: Set up Docker Buildx 18 | uses: docker/setup-buildx-action@v3 19 | with: 20 | version: latest 21 | - name: Login to GitHub Container Registry 22 | uses: docker/login-action@v3 23 | with: 24 | registry: ghcr.io 25 | username: ${{ github.actor }} 26 | password: ${{ secrets.GITHUB_TOKEN }} 27 | - name: Build and push Docker image 28 | run: | 29 | docker buildx bake \ 30 | -f docker-compose.yml \ 31 | -f .github/docker-compose.gha.yml \ 32 | --set *.platform=linux/arm64,linux/amd64,linux/arm/v7 \ 33 | --push pubsub 34 | -------------------------------------------------------------------------------- /.github/workflows/test.yml: -------------------------------------------------------------------------------- 1 | name: Test and Build 2 | 3 | on: 4 | workflow_dispatch: 5 | push: 6 | branches: 7 | - main 8 | pull_request: 9 | branches: 10 | - main 11 | 12 | jobs: 13 | test: 14 | runs-on: ubuntu-22.04 15 | steps: 16 | - name: Set up Go 17 | uses: actions/setup-go@v3 18 | with: 19 | go-version: 1.22 20 | - name: Checkout repo 21 | uses: actions/checkout@v3 22 | - name: Lint the codebase 23 | uses: golangci/golangci-lint-action@v3 24 | with: 25 | version: latest 26 | args: -E goimports -E godot --timeout 10m 27 | - name: Run tests 28 | run: | 29 | go test -v ./... -cover -race -coverprofile=coverage.out 30 | go tool cover -func=coverage.out -o=coverage.out 31 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.swp 2 | .DS_Store 3 | *.log 4 | .env 5 | ssh_data/ 6 | coverage.out 7 | -------------------------------------------------------------------------------- /.vscode/launch.json: -------------------------------------------------------------------------------- 1 | { 2 | // Use IntelliSense to learn about possible attributes. 3 | // Hover to view descriptions of existing attributes. 4 | // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 5 | "version": "0.2.0", 6 | "configurations": [ 7 | { 8 | "name": "Launch Package", 9 | "type": "go", 10 | "request": "launch", 11 | "mode": "auto", 12 | "program": "${workspaceFolder}/cmd/example/main.go", 13 | "cwd": "${workspaceFolder}", 14 | "buildFlags": [ 15 | "-race" 16 | ] 17 | } 18 | ] 19 | } -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | FROM --platform=$BUILDPLATFORM golang:1.23-alpine AS builder 2 | 3 | ENV CGO_ENABLED=0 4 | 5 | WORKDIR /usr/src/app 6 | 7 | COPY go.* . 8 | 9 | RUN --mount=type=cache,target=/go/pkg/,rw \ 10 | --mount=type=cache,target=/root/.cache/,rw \ 11 | go mod download 12 | 13 | WORKDIR /usr/src/app/cmd/example 14 | 15 | COPY ./cmd/example/go.* . 16 | 17 | RUN --mount=type=cache,target=/go/pkg/,rw \ 18 | --mount=type=cache,target=/root/.cache/,rw \ 19 | go mod download 20 | 21 | WORKDIR /usr/src/app 22 | 23 | COPY . . 24 | 25 | ARG TARGETOS 26 | ARG TARGETARCH 27 | 28 | ENV GOOS=${TARGETOS} GOARCH=${TARGETARCH} 29 | 30 | WORKDIR /usr/src/app/cmd/example 31 | 32 | RUN --mount=type=cache,target=/go/pkg/,rw \ 33 | --mount=type=cache,target=/root/.cache/,rw \ 34 | go build -o /usr/src/app/pubsub ./ 35 | 36 | FROM alpine 37 | 38 | RUN apk add --no-cache curl 39 | 40 | COPY --from=builder /usr/src/app/pubsub /pubsub 41 | 42 | ENTRYPOINT [ "/pubsub" ] 43 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 pico.sh LLC 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | fmt: 2 | go fmt ./... 3 | .PHONY: fmt 4 | 5 | lint: 6 | golangci-lint run -E goimports -E godot --timeout 10m 7 | .PHONY: lint 8 | 9 | test: 10 | go test -v ./... -cover -race -coverprofile=coverage.out 11 | go tool cover -func=coverage.out -o=coverage.out 12 | .PHONY: test 13 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # pubsub 2 | 3 | A generic pubsub implementation for Go. 4 | 5 | ```go 6 | package main 7 | 8 | import ( 9 | "bytes" 10 | "context" 11 | "fmt" 12 | "log/slog" 13 | 14 | "github.com/picosh/pubsub" 15 | ) 16 | 17 | func main() { 18 | ctx := context.TODO() 19 | logger := slog.Default() 20 | broker := pubsub.NewMulticast(logger) 21 | 22 | chann := []*pubsub.Channel{ 23 | pubsub.NewChannel("my-topic"), 24 | } 25 | 26 | go func() { 27 | writer := bytes.NewBufferString("my data") 28 | _ = broker.Pub(ctx, "pubID", writer, chann, false) 29 | }() 30 | 31 | reader := bytes.NewBufferString("") 32 | _ = broker.Sub(ctx, "subID", reader, chann, false) 33 | 34 | // result 35 | fmt.Println("data from pub:", reader) 36 | } 37 | ``` 38 | 39 | ## pubsub over ssh 40 | 41 | The simplest pubsub system for everyday automation needs. 42 | 43 | Using `wish` we can integrate our pubsub system into an SSH app. 44 | 45 | [![asciicast](https://asciinema.org/a/674287.svg)](https://asciinema.org/a/674287) 46 | 47 | ```bash 48 | # term 1 49 | mkdir ./ssh_data 50 | cat ~/.ssh/id_ed25519 ./ssh_data/authorized_keys 51 | go run ./cmd/example 52 | 53 | # term 2 54 | ssh -p 2222 localhost sub xyz 55 | 56 | # term 3 57 | echo "hello world" | ssh -p 2222 localhost pub xyz 58 | ``` 59 | -------------------------------------------------------------------------------- /broker.go: -------------------------------------------------------------------------------- 1 | package pubsub 2 | 3 | import ( 4 | "errors" 5 | "io" 6 | "iter" 7 | "log/slog" 8 | "sync" 9 | "time" 10 | 11 | "github.com/antoniomika/syncmap" 12 | ) 13 | 14 | /* 15 | Broker receives published messages and dispatches the message to the 16 | subscribing clients. An message contains a message topic that clients 17 | subscribe to and brokers use these subscription lists for determining the 18 | clients to receive the message. 19 | */ 20 | type Broker interface { 21 | GetChannels() iter.Seq2[string, *Channel] 22 | GetClients() iter.Seq2[string, *Client] 23 | Connect(*Client, []*Channel) (error, error) 24 | } 25 | 26 | type BaseBroker struct { 27 | Channels *syncmap.Map[string, *Channel] 28 | Logger *slog.Logger 29 | } 30 | 31 | func (b *BaseBroker) Cleanup() { 32 | toRemove := []string{} 33 | for _, channel := range b.GetChannels() { 34 | count := 0 35 | 36 | for range channel.GetClients() { 37 | count++ 38 | } 39 | 40 | if count == 0 { 41 | channel.Cleanup() 42 | toRemove = append(toRemove, channel.Topic) 43 | } 44 | } 45 | 46 | for _, channel := range toRemove { 47 | b.Channels.Delete(channel) 48 | } 49 | } 50 | 51 | func (b *BaseBroker) GetChannels() iter.Seq2[string, *Channel] { 52 | return b.Channels.Range 53 | } 54 | 55 | func (b *BaseBroker) GetClients() iter.Seq2[string, *Client] { 56 | return func(yield func(string, *Client) bool) { 57 | for _, channel := range b.GetChannels() { 58 | channel.Clients.Range(yield) 59 | } 60 | } 61 | } 62 | 63 | func (b *BaseBroker) Connect(client *Client, channels []*Channel) (error, error) { 64 | for _, channel := range channels { 65 | dataChannel := b.ensureChannel(channel) 66 | dataChannel.Clients.Store(client.ID, client) 67 | client.Channels.Store(dataChannel.Topic, dataChannel) 68 | defer func() { 69 | client.Channels.Delete(channel.Topic) 70 | dataChannel.Clients.Delete(client.ID) 71 | 72 | client.Cleanup() 73 | 74 | count := 0 75 | for _, cl := range dataChannel.GetClients() { 76 | if cl.Direction == ChannelDirectionInput || cl.Direction == ChannelDirectionInputOutput { 77 | count++ 78 | } 79 | } 80 | 81 | if count == 0 { 82 | for _, cl := range dataChannel.GetClients() { 83 | if !cl.KeepAlive { 84 | cl.Cleanup() 85 | } 86 | } 87 | } 88 | 89 | b.Cleanup() 90 | }() 91 | } 92 | 93 | var ( 94 | inputErr error 95 | outputErr error 96 | wg sync.WaitGroup 97 | ) 98 | 99 | // Pub 100 | if client.Direction == ChannelDirectionInput || client.Direction == ChannelDirectionInputOutput { 101 | wg.Add(1) 102 | go func() { 103 | defer wg.Done() 104 | for { 105 | data := make([]byte, 32*1024) 106 | n, err := client.ReadWriter.Read(data) 107 | data = data[:n] 108 | 109 | channelMessage := ChannelMessage{ 110 | Data: data, 111 | ClientID: client.ID, 112 | Direction: ChannelDirectionInput, 113 | } 114 | 115 | if client.BlockWrite { 116 | mainLoop: 117 | for { 118 | count := 0 119 | for _, channel := range client.GetChannels() { 120 | for _, chanClient := range channel.GetClients() { 121 | if chanClient.Direction == ChannelDirectionOutput || chanClient.Direction == ChannelDirectionInputOutput { 122 | count++ 123 | } 124 | } 125 | } 126 | 127 | if count > 0 { 128 | break mainLoop 129 | } 130 | 131 | select { 132 | case <-client.Done: 133 | break mainLoop 134 | case <-time.After(1 * time.Millisecond): 135 | continue 136 | } 137 | } 138 | } 139 | 140 | var sendwg sync.WaitGroup 141 | 142 | for _, channel := range client.GetChannels() { 143 | sendwg.Add(1) 144 | go func() { 145 | defer sendwg.Done() 146 | select { 147 | case channel.Data <- channelMessage: 148 | case <-client.Done: 149 | case <-channel.Done: 150 | } 151 | }() 152 | } 153 | 154 | sendwg.Wait() 155 | 156 | if err != nil { 157 | if errors.Is(err, io.EOF) { 158 | return 159 | } 160 | inputErr = err 161 | return 162 | } 163 | } 164 | }() 165 | } 166 | 167 | // Sub 168 | if client.Direction == ChannelDirectionOutput || client.Direction == ChannelDirectionInputOutput { 169 | wg.Add(1) 170 | go func() { 171 | defer wg.Done() 172 | mainLoop: 173 | for { 174 | select { 175 | case data, ok := <-client.Data: 176 | _, err := client.ReadWriter.Write(data.Data) 177 | if err != nil { 178 | outputErr = err 179 | break mainLoop 180 | } 181 | 182 | if !ok { 183 | break mainLoop 184 | } 185 | case <-client.Done: 186 | break mainLoop 187 | } 188 | } 189 | }() 190 | } 191 | 192 | wg.Wait() 193 | 194 | return inputErr, outputErr 195 | } 196 | 197 | func (b *BaseBroker) ensureChannel(channel *Channel) *Channel { 198 | dataChannel, _ := b.Channels.LoadOrStore(channel.Topic, channel) 199 | dataChannel.Handle() 200 | return dataChannel 201 | } 202 | 203 | var _ Broker = (*BaseBroker)(nil) 204 | -------------------------------------------------------------------------------- /channel.go: -------------------------------------------------------------------------------- 1 | package pubsub 2 | 3 | import ( 4 | "iter" 5 | "sync" 6 | 7 | "github.com/antoniomika/syncmap" 8 | ) 9 | 10 | type ChannelDirection int 11 | 12 | func (d ChannelDirection) String() string { 13 | return [...]string{"input", "output", "inputoutput"}[d] 14 | } 15 | 16 | const ( 17 | ChannelDirectionInput ChannelDirection = iota 18 | ChannelDirectionOutput 19 | ChannelDirectionInputOutput 20 | ) 21 | 22 | type ChannelAction int 23 | 24 | func (d ChannelAction) String() string { 25 | return [...]string{"data", "close"}[d] 26 | } 27 | 28 | const ( 29 | ChannelActionData = iota 30 | ChannelActionClose 31 | ) 32 | 33 | type ChannelMessage struct { 34 | Data []byte 35 | ClientID string 36 | Direction ChannelDirection 37 | Action ChannelAction 38 | } 39 | 40 | func NewChannel(topic string) *Channel { 41 | return &Channel{ 42 | Topic: topic, 43 | Done: make(chan struct{}), 44 | Data: make(chan ChannelMessage), 45 | Clients: syncmap.New[string, *Client](), 46 | } 47 | } 48 | 49 | /* 50 | Channel is a container for a topic. It holds the list of clients and 51 | a data channel to receive a message. 52 | */ 53 | type Channel struct { 54 | Topic string 55 | Done chan struct{} 56 | Data chan ChannelMessage 57 | Clients *syncmap.Map[string, *Client] 58 | handleOnce sync.Once 59 | cleanupOnce sync.Once 60 | } 61 | 62 | func (c *Channel) GetClients() iter.Seq2[string, *Client] { 63 | return c.Clients.Range 64 | } 65 | 66 | func (c *Channel) Cleanup() { 67 | c.cleanupOnce.Do(func() { 68 | close(c.Done) 69 | }) 70 | } 71 | 72 | func (c *Channel) Handle() { 73 | c.handleOnce.Do(func() { 74 | go func() { 75 | defer func() { 76 | for _, client := range c.GetClients() { 77 | client.Cleanup() 78 | } 79 | }() 80 | 81 | for { 82 | select { 83 | case <-c.Done: 84 | return 85 | case data, ok := <-c.Data: 86 | var wg sync.WaitGroup 87 | for _, client := range c.GetClients() { 88 | if client.Direction == ChannelDirectionInput || (client.ID == data.ClientID && !client.Replay) { 89 | continue 90 | } 91 | 92 | wg.Add(1) 93 | go func() { 94 | defer wg.Done() 95 | if !ok { 96 | client.onceData.Do(func() { 97 | close(client.Data) 98 | }) 99 | return 100 | } 101 | 102 | select { 103 | case client.Data <- data: 104 | case <-client.Done: 105 | case <-c.Done: 106 | } 107 | }() 108 | } 109 | wg.Wait() 110 | } 111 | } 112 | }() 113 | }) 114 | } 115 | -------------------------------------------------------------------------------- /client.go: -------------------------------------------------------------------------------- 1 | package pubsub 2 | 3 | import ( 4 | "io" 5 | "iter" 6 | "sync" 7 | 8 | "github.com/antoniomika/syncmap" 9 | ) 10 | 11 | func NewClient(ID string, rw io.ReadWriter, direction ChannelDirection, blockWrite, replay, keepAlive bool) *Client { 12 | return &Client{ 13 | ID: ID, 14 | ReadWriter: rw, 15 | Direction: direction, 16 | Channels: syncmap.New[string, *Channel](), 17 | Done: make(chan struct{}), 18 | Data: make(chan ChannelMessage), 19 | Replay: replay, 20 | BlockWrite: blockWrite, 21 | KeepAlive: keepAlive, 22 | } 23 | } 24 | 25 | /* 26 | Client is the container for holding state between multiple devices. A 27 | client has a direction (input, output, inputout) as well as a way to 28 | send data to all the associated channels. 29 | */ 30 | type Client struct { 31 | ID string 32 | ReadWriter io.ReadWriter 33 | Channels *syncmap.Map[string, *Channel] 34 | Direction ChannelDirection 35 | Done chan struct{} 36 | Data chan ChannelMessage 37 | Replay bool 38 | BlockWrite bool 39 | KeepAlive bool 40 | once sync.Once 41 | onceData sync.Once 42 | } 43 | 44 | func (c *Client) GetChannels() iter.Seq2[string, *Channel] { 45 | return c.Channels.Range 46 | } 47 | 48 | func (c *Client) Cleanup() { 49 | c.once.Do(func() { 50 | close(c.Done) 51 | }) 52 | } 53 | -------------------------------------------------------------------------------- /cmd/basic/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "bytes" 5 | "context" 6 | "fmt" 7 | "log/slog" 8 | 9 | "github.com/picosh/pubsub" 10 | ) 11 | 12 | func main() { 13 | ctx := context.TODO() 14 | logger := slog.Default() 15 | broker := pubsub.NewMulticast(logger) 16 | 17 | chann := []*pubsub.Channel{ 18 | pubsub.NewChannel("my-topic"), 19 | } 20 | 21 | go func() { 22 | writer := bytes.NewBufferString("my data") 23 | _ = broker.Pub(ctx, "pubID", writer, chann, false) 24 | }() 25 | 26 | reader := bytes.NewBufferString("") 27 | _ = broker.Sub(ctx, "subID", reader, chann, false) 28 | 29 | // result 30 | fmt.Println("data from pub:", reader) 31 | } 32 | -------------------------------------------------------------------------------- /cmd/example/go.mod: -------------------------------------------------------------------------------- 1 | module github.com/picosh/pubsub/cmd/example 2 | 3 | go 1.23.1 4 | 5 | replace github.com/picosh/pubsub => ../../ 6 | 7 | require ( 8 | github.com/charmbracelet/ssh v0.0.0-20240725163421-eb71b85b27aa 9 | github.com/charmbracelet/wish v1.4.3 10 | github.com/google/uuid v1.6.0 11 | github.com/picosh/pubsub v0.0.0-20241003193028-85a65ebffeb6 12 | ) 13 | 14 | require ( 15 | github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be // indirect 16 | github.com/antoniomika/syncmap v1.0.0 // indirect 17 | github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect 18 | github.com/charmbracelet/bubbletea v1.1.1 // indirect 19 | github.com/charmbracelet/keygen v0.5.1 // indirect 20 | github.com/charmbracelet/lipgloss v0.13.0 // indirect 21 | github.com/charmbracelet/log v0.4.0 // indirect 22 | github.com/charmbracelet/x/ansi v0.3.2 // indirect 23 | github.com/charmbracelet/x/conpty v0.1.0 // indirect 24 | github.com/charmbracelet/x/errors v0.0.0-20241007193646-7cc13b2883e3 // indirect 25 | github.com/charmbracelet/x/term v0.2.0 // indirect 26 | github.com/charmbracelet/x/termios v0.1.0 // indirect 27 | github.com/creack/pty v1.1.23 // indirect 28 | github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect 29 | github.com/go-logfmt/logfmt v0.6.0 // indirect 30 | github.com/lucasb-eyer/go-colorful v1.2.0 // indirect 31 | github.com/mattn/go-isatty v0.0.20 // indirect 32 | github.com/mattn/go-localereader v0.0.1 // indirect 33 | github.com/mattn/go-runewidth v0.0.16 // indirect 34 | github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect 35 | github.com/muesli/cancelreader v0.2.2 // indirect 36 | github.com/muesli/termenv v0.15.3-0.20240509142007-81b8f94111d5 // indirect 37 | github.com/rivo/uniseg v0.4.7 // indirect 38 | golang.org/x/crypto v0.28.0 // indirect 39 | golang.org/x/exp v0.0.0-20241004190924-225e2abe05e6 // indirect 40 | golang.org/x/sync v0.8.0 // indirect 41 | golang.org/x/sys v0.26.0 // indirect 42 | golang.org/x/text v0.19.0 // indirect 43 | ) 44 | -------------------------------------------------------------------------------- /cmd/example/go.sum: -------------------------------------------------------------------------------- 1 | github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be h1:9AeTilPcZAjCFIImctFaOjnTIavg87rW78vTPkQqLI8= 2 | github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be/go.mod h1:ySMOLuWl6zY27l47sB3qLNK6tF2fkHG55UZxx8oIVo4= 3 | github.com/antoniomika/syncmap v1.0.0 h1:iFSfbQFQOvHZILFZF+hqWosO0no+W9+uF4y2VEyMKWU= 4 | github.com/antoniomika/syncmap v1.0.0/go.mod h1:fK2829foEYnO4riNfyUn0SHQZt4ue3DStYjGU+sJj38= 5 | github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k= 6 | github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8= 7 | github.com/charmbracelet/bubbletea v1.1.1 h1:KJ2/DnmpfqFtDNVTvYZ6zpPFL9iRCRr0qqKOCvppbPY= 8 | github.com/charmbracelet/bubbletea v1.1.1/go.mod h1:9Ogk0HrdbHolIKHdjfFpyXJmiCzGwy+FesYkZr7hYU4= 9 | github.com/charmbracelet/keygen v0.5.1 h1:zBkkYPtmKDVTw+cwUyY6ZwGDhRxXkEp0Oxs9sqMLqxI= 10 | github.com/charmbracelet/keygen v0.5.1/go.mod h1:zznJVmK/GWB6dAtjluqn2qsttiCBhA5MZSiwb80fcHw= 11 | github.com/charmbracelet/lipgloss v0.13.0 h1:4X3PPeoWEDCMvzDvGmTajSyYPcZM4+y8sCA/SsA3cjw= 12 | github.com/charmbracelet/lipgloss v0.13.0/go.mod h1:nw4zy0SBX/F/eAO1cWdcvy6qnkDUxr8Lw7dvFrAIbbY= 13 | github.com/charmbracelet/log v0.4.0 h1:G9bQAcx8rWA2T3pWvx7YtPTPwgqpk7D68BX21IRW8ZM= 14 | github.com/charmbracelet/log v0.4.0/go.mod h1:63bXt/djrizTec0l11H20t8FDSvA4CRZJ1KH22MdptM= 15 | github.com/charmbracelet/ssh v0.0.0-20240725163421-eb71b85b27aa h1:6rePgmsJguB6Z7Y55stsEVDlWFJoUpQvOX4mdnBjgx4= 16 | github.com/charmbracelet/ssh v0.0.0-20240725163421-eb71b85b27aa/go.mod h1:LmMZag2g7ILMmWtDmU7dIlctUopwmb73KpPzj0ip1uk= 17 | github.com/charmbracelet/wish v1.4.3 h1:7FvNLoPGqiT7EdjQP4+XuvM1Hrnx9DyknilbD+Okx1s= 18 | github.com/charmbracelet/wish v1.4.3/go.mod h1:hVgmhwhd52fLmO6m5AkREUMZYqQ0qmIJQDMe3HsNPmU= 19 | github.com/charmbracelet/x/ansi v0.3.2 h1:wsEwgAN+C9U06l9dCVMX0/L3x7ptvY1qmjMwyfE6USY= 20 | github.com/charmbracelet/x/ansi v0.3.2/go.mod h1:dk73KoMTT5AX5BsX0KrqhsTqAnhZZoCBjs7dGWp4Ktw= 21 | github.com/charmbracelet/x/conpty v0.1.0 h1:4zc8KaIcbiL4mghEON8D72agYtSeIgq8FSThSPQIb+U= 22 | github.com/charmbracelet/x/conpty v0.1.0/go.mod h1:rMFsDJoDwVmiYM10aD4bH2XiRgwI7NYJtQgl5yskjEQ= 23 | github.com/charmbracelet/x/errors v0.0.0-20241007193646-7cc13b2883e3 h1:nsBhzPXBqeXEGZ9ztveSIPdf790BcDikbaEH3vMglH4= 24 | github.com/charmbracelet/x/errors v0.0.0-20241007193646-7cc13b2883e3/go.mod h1:2P0UgXMEa6TsToMSuFqKFQR+fZTO9CNGUNokkPatT/0= 25 | github.com/charmbracelet/x/term v0.2.0 h1:cNB9Ot9q8I711MyZ7myUR5HFWL/lc3OpU8jZ4hwm0x0= 26 | github.com/charmbracelet/x/term v0.2.0/go.mod h1:GVxgxAbjUrmpvIINHIQnJJKpMlHiZ4cktEQCN6GWyF0= 27 | github.com/charmbracelet/x/termios v0.1.0 h1:y4rjAHeFksBAfGbkRDmVinMg7x7DELIGAFbdNvxg97k= 28 | github.com/charmbracelet/x/termios v0.1.0/go.mod h1:H/EVv/KRnrYjz+fCYa9bsKdqF3S8ouDK0AZEbG7r+/U= 29 | github.com/creack/pty v1.1.23 h1:4M6+isWdcStXEf15G/RbrMPOQj1dZ7HPZCGwE4kOeP0= 30 | github.com/creack/pty v1.1.23/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE= 31 | github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= 32 | github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 33 | github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4= 34 | github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f/go.mod h1:vw97MGsxSvLiUE2X8qFplwetxpGLQrlU1Q9AUEIzCaM= 35 | github.com/go-logfmt/logfmt v0.6.0 h1:wGYYu3uicYdqXVgoYbvnkrPVXkuLM1p1ifugDMEdRi4= 36 | github.com/go-logfmt/logfmt v0.6.0/go.mod h1:WYhtIu8zTZfxdn5+rREduYbwxfcBr/Vr6KEVveWlfTs= 37 | github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= 38 | github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= 39 | github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY= 40 | github.com/lucasb-eyer/go-colorful v1.2.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= 41 | github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= 42 | github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= 43 | github.com/mattn/go-localereader v0.0.1 h1:ygSAOl7ZXTx4RdPYinUpg6W99U8jWvWi9Ye2JC/oIi4= 44 | github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+EiG4R1k4Cjx5p88= 45 | github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc= 46 | github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= 47 | github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI= 48 | github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6/go.mod h1:CJlz5H+gyd6CUWT45Oy4q24RdLyn7Md9Vj2/ldJBSIo= 49 | github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA= 50 | github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo= 51 | github.com/muesli/termenv v0.15.3-0.20240509142007-81b8f94111d5 h1:NiONcKK0EV5gUZcnCiPMORaZA0eBDc+Fgepl9xl4lZ8= 52 | github.com/muesli/termenv v0.15.3-0.20240509142007-81b8f94111d5/go.mod h1:hxSnBBYLK21Vtq/PHd0S2FYCxBXzBua8ov5s1RobyRQ= 53 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 54 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 55 | github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= 56 | github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ= 57 | github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= 58 | github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= 59 | github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= 60 | golang.org/x/crypto v0.28.0 h1:GBDwsMXVQi34v5CCYUm2jkJvu4cbtru2U4TN2PSyQnw= 61 | golang.org/x/crypto v0.28.0/go.mod h1:rmgy+3RHxRZMyY0jjAJShp2zgEdOqj2AO7U0pYmeQ7U= 62 | golang.org/x/exp v0.0.0-20241004190924-225e2abe05e6 h1:1wqE9dj9NpSm04INVsJhhEUzhuDVjbcyKH91sVyPATw= 63 | golang.org/x/exp v0.0.0-20241004190924-225e2abe05e6/go.mod h1:NQtJDoLvd6faHhE7m4T/1IY708gDefGGjR/iUW8yQQ8= 64 | golang.org/x/sync v0.8.0 h1:3NFvSEYkUoMifnESzZl15y791HH1qU2xm6eCJU5ZPXQ= 65 | golang.org/x/sync v0.8.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= 66 | golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 67 | golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 68 | golang.org/x/sys v0.26.0 h1:KHjCJyddX0LoSTb3J+vWpupP9p0oznkqVk/IfjymZbo= 69 | golang.org/x/sys v0.26.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= 70 | golang.org/x/term v0.25.0 h1:WtHI/ltw4NvSUig5KARz9h521QvRC8RmF/cuYqifU24= 71 | golang.org/x/term v0.25.0/go.mod h1:RPyXicDX+6vLxogjjRxjgD2TKtmAO6NZBsBRfrOLu7M= 72 | golang.org/x/text v0.19.0 h1:kTxAhCbGbxhK0IwgSKiMO5awPoDQ0RpfiVYBfK860YM= 73 | golang.org/x/text v0.19.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY= 74 | gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= 75 | gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= 76 | -------------------------------------------------------------------------------- /cmd/example/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "fmt" 7 | "log/slog" 8 | "os" 9 | "os/signal" 10 | "runtime" 11 | "strings" 12 | "syscall" 13 | "time" 14 | 15 | "github.com/charmbracelet/ssh" 16 | "github.com/charmbracelet/wish" 17 | "github.com/google/uuid" 18 | "github.com/picosh/pubsub" 19 | ) 20 | 21 | func GetEnv(key string, defaultVal string) string { 22 | if value, exists := os.LookupEnv(key); exists { 23 | return value 24 | } 25 | return defaultVal 26 | } 27 | 28 | func PubSubMiddleware(broker pubsub.PubSub, logger *slog.Logger) wish.Middleware { 29 | return func(next ssh.Handler) ssh.Handler { 30 | return func(sesh ssh.Session) { 31 | args := sesh.Command() 32 | if len(args) < 2 { 33 | wish.Println(sesh, "USAGE: ssh send.pico.sh (sub|pub|pipe) {topic}") 34 | next(sesh) 35 | return 36 | } 37 | 38 | cmd := strings.TrimSpace(args[0]) 39 | topicsRaw := args[1] 40 | 41 | topics := strings.Split(topicsRaw, ",") 42 | 43 | logger := logger.With( 44 | "cmd", cmd, 45 | "topics", topics, 46 | ) 47 | 48 | logger.Info("running cli") 49 | 50 | if cmd == "help" { 51 | wish.Println(sesh, "USAGE: ssh send.pico.sh (sub|pub|pipe) {topic}") 52 | } else if cmd == "sub" { 53 | var chans []*pubsub.Channel 54 | 55 | for _, topic := range topics { 56 | chans = append(chans, pubsub.NewChannel(topic)) 57 | } 58 | 59 | clientID := uuid.NewString() 60 | 61 | err := errors.Join(broker.Sub(sesh.Context(), clientID, sesh, chans, strings.ToLower(args[len(args)-1]) == "keepalive")) 62 | if err != nil { 63 | logger.Error("error during pub", slog.Any("error", err), slog.String("client", clientID)) 64 | } 65 | } else if cmd == "pub" { 66 | var chans []*pubsub.Channel 67 | 68 | for _, topic := range topics { 69 | chans = append(chans, pubsub.NewChannel(topic)) 70 | } 71 | 72 | clientID := uuid.NewString() 73 | 74 | err := errors.Join(broker.Pub(sesh.Context(), clientID, sesh, chans, strings.ToLower(args[len(args)-1]) == "blockwrite")) 75 | if err != nil { 76 | logger.Error("error during pub", slog.Any("error", err), slog.String("client", clientID)) 77 | } 78 | } else if cmd == "pipe" { 79 | var chans []*pubsub.Channel 80 | 81 | for _, topics := range topics { 82 | chans = append(chans, pubsub.NewChannel(topics)) 83 | } 84 | 85 | clientID := uuid.NewString() 86 | 87 | err := errors.Join(broker.Pipe(sesh.Context(), clientID, sesh, chans, strings.ToLower(args[len(args)-1]) == "replay")) 88 | if err != nil { 89 | logger.Error( 90 | "pipe error", 91 | slog.Any("error", err), 92 | slog.String("pipeClient", clientID), 93 | ) 94 | } 95 | } else { 96 | wish.Println(sesh, "USAGE: ssh send.pico.sh (sub|pub|pipe) {topic}") 97 | } 98 | 99 | next(sesh) 100 | } 101 | } 102 | } 103 | 104 | func main() { 105 | logger := slog.Default() 106 | host := GetEnv("SSH_HOST", "0.0.0.0") 107 | port := GetEnv("SSH_PORT", "2222") 108 | keyPath := GetEnv("SSH_AUTHORIZED_KEYS", "./ssh_data/authorized_keys") 109 | broker := pubsub.NewMulticast(logger) 110 | 111 | s, err := wish.NewServer( 112 | ssh.NoPty(), 113 | wish.WithAddress(fmt.Sprintf("%s:%s", host, port)), 114 | wish.WithHostKeyPath("ssh_data/term_info_ed25519"), 115 | wish.WithAuthorizedKeys(keyPath), 116 | wish.WithMiddleware( 117 | PubSubMiddleware(broker, logger), 118 | ), 119 | ) 120 | if err != nil { 121 | logger.Error(err.Error()) 122 | return 123 | } 124 | 125 | done := make(chan os.Signal, 1) 126 | signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM) 127 | logger.Info( 128 | "starting SSH server", 129 | "host", host, 130 | "port", port, 131 | ) 132 | go func() { 133 | if err = s.ListenAndServe(); err != nil { 134 | logger.Error(err.Error()) 135 | } 136 | }() 137 | 138 | go func() { 139 | for { 140 | slog.Info("Debug Info", slog.Int("goroutines", runtime.NumGoroutine())) 141 | select { 142 | case <-time.After(5 * time.Second): 143 | for _, channel := range broker.GetChannels() { 144 | slog.Info("channel online", slog.Any("channel topic", channel.Topic)) 145 | for _, client := range channel.GetClients() { 146 | slog.Info("client online", slog.Any("channel topic", channel.Topic), slog.Any("client", client.ID), slog.String("direction", client.Direction.String())) 147 | } 148 | } 149 | case <-done: 150 | return 151 | } 152 | } 153 | }() 154 | 155 | <-done 156 | logger.Info("stopping SSH server") 157 | ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) 158 | defer func() { 159 | cancel() 160 | }() 161 | if err := s.Shutdown(ctx); err != nil { 162 | logger.Error(err.Error()) 163 | } 164 | } 165 | -------------------------------------------------------------------------------- /docker-compose.yml: -------------------------------------------------------------------------------- 1 | services: 2 | pubsub: 3 | build: 4 | context: . 5 | image: ghcr.io/picosh/pubsub:latest 6 | restart: always 7 | volumes: 8 | - ./ssh_data:/ssh_data 9 | ports: 10 | - 2222:2222 11 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/picosh/pubsub 2 | 3 | go 1.23.1 4 | 5 | require github.com/antoniomika/syncmap v1.0.0 6 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/antoniomika/syncmap v1.0.0 h1:iFSfbQFQOvHZILFZF+hqWosO0no+W9+uF4y2VEyMKWU= 2 | github.com/antoniomika/syncmap v1.0.0/go.mod h1:fK2829foEYnO4riNfyUn0SHQZt4ue3DStYjGU+sJj38= 3 | -------------------------------------------------------------------------------- /multicast.go: -------------------------------------------------------------------------------- 1 | package pubsub 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "io" 7 | "iter" 8 | "log/slog" 9 | 10 | "github.com/antoniomika/syncmap" 11 | ) 12 | 13 | /* 14 | Multicast is a flexible, bidirectional broker. 15 | 16 | It provides the most pure version of our PubSub interface which lets 17 | end-developers build one-to-many connections between publishers and 18 | subscribers and vice versa. 19 | 20 | It doesn't provide any topic filtering capabilities and is only 21 | concerned with sending data to and from an `io.ReadWriter` via our 22 | channels. 23 | */ 24 | type Multicast struct { 25 | Broker 26 | Logger *slog.Logger 27 | } 28 | 29 | func NewMulticast(logger *slog.Logger) *Multicast { 30 | return &Multicast{ 31 | Logger: logger, 32 | Broker: &BaseBroker{ 33 | Channels: syncmap.New[string, *Channel](), 34 | Logger: logger.With(slog.Bool("broker", true)), 35 | }, 36 | } 37 | } 38 | 39 | func (p *Multicast) getClients(direction ChannelDirection) iter.Seq2[string, *Client] { 40 | return func(yield func(string, *Client) bool) { 41 | for clientID, client := range p.GetClients() { 42 | if client.Direction == direction { 43 | yield(clientID, client) 44 | } 45 | } 46 | } 47 | } 48 | 49 | func (p *Multicast) GetPipes() iter.Seq2[string, *Client] { 50 | return p.getClients(ChannelDirectionInputOutput) 51 | } 52 | 53 | func (p *Multicast) GetPubs() iter.Seq2[string, *Client] { 54 | return p.getClients(ChannelDirectionInput) 55 | } 56 | 57 | func (p *Multicast) GetSubs() iter.Seq2[string, *Client] { 58 | return p.getClients(ChannelDirectionOutput) 59 | } 60 | 61 | func (p *Multicast) connect(ctx context.Context, ID string, rw io.ReadWriter, channels []*Channel, direction ChannelDirection, blockWrite bool, replay, keepAlive bool) (error, error) { 62 | client := NewClient(ID, rw, direction, blockWrite, replay, keepAlive) 63 | 64 | go func() { 65 | <-ctx.Done() 66 | client.Cleanup() 67 | }() 68 | 69 | return p.Connect(client, channels) 70 | } 71 | 72 | func (p *Multicast) Pipe(ctx context.Context, ID string, rw io.ReadWriter, channels []*Channel, replay bool) (error, error) { 73 | return p.connect(ctx, ID, rw, channels, ChannelDirectionInputOutput, false, replay, false) 74 | } 75 | 76 | func (p *Multicast) Pub(ctx context.Context, ID string, rw io.ReadWriter, channels []*Channel, blockWrite bool) error { 77 | return errors.Join(p.connect(ctx, ID, rw, channels, ChannelDirectionInput, blockWrite, false, false)) 78 | } 79 | 80 | func (p *Multicast) Sub(ctx context.Context, ID string, rw io.ReadWriter, channels []*Channel, keepAlive bool) error { 81 | return errors.Join(p.connect(ctx, ID, rw, channels, ChannelDirectionOutput, false, false, keepAlive)) 82 | } 83 | 84 | var _ PubSub = (*Multicast)(nil) 85 | -------------------------------------------------------------------------------- /multicast_test.go: -------------------------------------------------------------------------------- 1 | package pubsub 2 | 3 | import ( 4 | "bytes" 5 | "context" 6 | "fmt" 7 | "log/slog" 8 | "sync" 9 | "testing" 10 | ) 11 | 12 | type Buffer struct { 13 | b bytes.Buffer 14 | m sync.Mutex 15 | } 16 | 17 | func (b *Buffer) Read(p []byte) (n int, err error) { 18 | b.m.Lock() 19 | defer b.m.Unlock() 20 | return b.b.Read(p) 21 | } 22 | func (b *Buffer) Write(p []byte) (n int, err error) { 23 | b.m.Lock() 24 | defer b.m.Unlock() 25 | return b.b.Write(p) 26 | } 27 | func (b *Buffer) String() string { 28 | b.m.Lock() 29 | defer b.m.Unlock() 30 | return b.b.String() 31 | } 32 | 33 | func TestMulticastSubBlock(t *testing.T) { 34 | orderActual := "" 35 | orderExpected := "sub-pub-" 36 | actual := new(Buffer) 37 | expected := "some test data" 38 | name := "test-channel" 39 | syncer := make(chan int) 40 | 41 | cast := NewMulticast(slog.Default()) 42 | 43 | var wg sync.WaitGroup 44 | wg.Add(2) 45 | 46 | channel := NewChannel(name) 47 | 48 | go func() { 49 | orderActual += "sub-" 50 | syncer <- 0 51 | fmt.Println(cast.Sub(context.TODO(), "1", actual, []*Channel{channel}, false)) 52 | wg.Done() 53 | }() 54 | 55 | <-syncer 56 | 57 | go func() { 58 | orderActual += "pub-" 59 | fmt.Println(cast.Pub(context.TODO(), "2", &Buffer{b: *bytes.NewBufferString(expected)}, []*Channel{channel}, true)) 60 | wg.Done() 61 | }() 62 | 63 | wg.Wait() 64 | 65 | if orderActual != orderExpected { 66 | t.Fatalf("\norderActual:(%s)\norderExpected:(%s)", orderActual, orderExpected) 67 | } 68 | if actual.String() != expected { 69 | t.Fatalf("\nactual:(%s)\nexpected:(%s)", actual, expected) 70 | } 71 | } 72 | 73 | func TestMulticastPubBlock(t *testing.T) { 74 | orderActual := "" 75 | orderExpected := "pub-sub-" 76 | actual := new(Buffer) 77 | expected := "some test data" 78 | name := "test-channel" 79 | syncer := make(chan int) 80 | 81 | cast := NewMulticast(slog.Default()) 82 | 83 | var wg sync.WaitGroup 84 | wg.Add(2) 85 | 86 | channel := NewChannel(name) 87 | 88 | go func() { 89 | orderActual += "pub-" 90 | syncer <- 0 91 | fmt.Println(cast.Pub(context.TODO(), "1", &Buffer{b: *bytes.NewBufferString(expected)}, []*Channel{channel}, true)) 92 | wg.Done() 93 | }() 94 | 95 | <-syncer 96 | 97 | go func() { 98 | orderActual += "sub-" 99 | wg.Done() 100 | fmt.Println(cast.Sub(context.TODO(), "2", actual, []*Channel{channel}, false)) 101 | }() 102 | 103 | wg.Wait() 104 | 105 | if orderActual != orderExpected { 106 | t.Fatalf("\norderActual:(%s)\norderExpected:(%s)", orderActual, orderExpected) 107 | } 108 | if actual.String() != expected { 109 | t.Fatalf("\nactual:(%s)\nexpected:(%s)", actual, expected) 110 | } 111 | } 112 | 113 | func TestMulticastMultSubs(t *testing.T) { 114 | orderActual := "" 115 | orderExpected := "sub-sub-pub-" 116 | actual := new(Buffer) 117 | actualOther := new(Buffer) 118 | expected := "some test data" 119 | name := "test-channel" 120 | syncer := make(chan int) 121 | 122 | cast := NewMulticast(slog.Default()) 123 | 124 | var wg sync.WaitGroup 125 | wg.Add(3) 126 | 127 | channel := NewChannel(name) 128 | 129 | go func() { 130 | orderActual += "sub-" 131 | syncer <- 0 132 | fmt.Println(cast.Sub(context.TODO(), "1", actual, []*Channel{channel}, false)) 133 | wg.Done() 134 | }() 135 | 136 | <-syncer 137 | 138 | go func() { 139 | orderActual += "sub-" 140 | syncer <- 0 141 | fmt.Println(cast.Sub(context.TODO(), "2", actualOther, []*Channel{channel}, false)) 142 | wg.Done() 143 | }() 144 | 145 | <-syncer 146 | 147 | go func() { 148 | orderActual += "pub-" 149 | fmt.Println(cast.Pub(context.TODO(), "3", &Buffer{b: *bytes.NewBufferString(expected)}, []*Channel{channel}, true)) 150 | wg.Done() 151 | }() 152 | 153 | wg.Wait() 154 | 155 | if orderActual != orderExpected { 156 | t.Fatalf("\norderActual:(%s)\norderExpected:(%s)", orderActual, orderExpected) 157 | } 158 | if actual.String() != expected { 159 | t.Fatalf("\nactual:(%s)\nexpected:(%s)", actual, expected) 160 | } 161 | if actualOther.String() != expected { 162 | t.Fatalf("\nactual:(%s)\nexpected:(%s)", actualOther, expected) 163 | } 164 | } 165 | -------------------------------------------------------------------------------- /pubsub.go: -------------------------------------------------------------------------------- 1 | package pubsub 2 | 3 | import ( 4 | "context" 5 | "io" 6 | "iter" 7 | ) 8 | 9 | /* 10 | PubSub is our take on a basic publisher and subscriber interface. 11 | 12 | It has a few notable requirements: 13 | - Each operation must accept an array of channels 14 | - A way to send, receive, and stream data between clients 15 | 16 | PubSub also inherits the properties of a Broker. 17 | */ 18 | type PubSub interface { 19 | Broker 20 | GetPubs() iter.Seq2[string, *Client] 21 | GetSubs() iter.Seq2[string, *Client] 22 | GetPipes() iter.Seq2[string, *Client] 23 | Pipe(ctx context.Context, ID string, rw io.ReadWriter, channels []*Channel, replay bool) (error, error) 24 | Sub(ctx context.Context, ID string, rw io.ReadWriter, channels []*Channel, keepAlive bool) error 25 | Pub(ctx context.Context, ID string, rw io.ReadWriter, channels []*Channel, blockWrite bool) error 26 | } 27 | --------------------------------------------------------------------------------