├── .circleci
    └── config.yml
├── .gitignore
├── .golangci.yml
├── LICENSE
├── Makefile
├── README.md
├── client.go
├── client_connect_error.go
├── client_connect_error_test.go
├── client_factory.go
├── client_factory_test.go
├── client_not_found_error.go
├── client_not_found_error_test.go
├── client_options.go
├── client_options_test.go
├── client_ping_error.go
├── client_ping_error_test.go
├── client_receive_error.go
├── client_receive_error_test.go
├── client_repeat_connect_error.go
├── client_repeat_connect_error_test.go
├── client_send_buffer_overflow_error.go
├── client_send_buffer_overflow_error_test.go
├── client_send_error.go
├── client_send_error_test.go
├── client_store.go
├── client_store_channels_shard.go
├── client_store_clients_shard.go
├── client_store_options.go
├── client_store_options_test.go
├── client_store_test.go
├── client_test.go
├── connection_closed_error.go
├── connection_closed_error_test.go
├── doc.go
├── examples
    ├── bench
    │   ├── client
    │   │   └── main.go
    │   └── server
    │   │   └── main.go
    └── demo
    │   ├── README.md
    │   ├── index.html
    │   └── main.go
├── go.mod
├── go.sum
├── gobwas_connection.go
├── gobwas_upgrader.go
├── gobwas_upgrader_options.go
├── gobwas_upgrader_options_test.go
├── gorilla_connection.go
├── gorilla_upgrader.go
├── gorilla_upgrader_options.go
├── gorilla_upgrader_options_test.go
├── hub.go
├── hub_options.go
├── hub_options_test.go
├── hub_subscription_channel_required_error.go
├── hub_subscription_channel_required_error_test.go
├── hub_test.go
├── logrus_logger.go
├── logrus_options.go
├── logrus_options_test.go
├── message.go
├── message_test.go
├── mock
    ├── client_factory.go
    └── hub.go
├── satori_uuid_generator.go
├── uuid.go
└── uuid_test.go


/.circleci/config.yml:
--------------------------------------------------------------------------------
 1 | version: 2.1
 2 | 
 3 | orbs:
 4 |     codecov: codecov/codecov@1.0.2
 5 | 
 6 | jobs:
 7 |     build:
 8 |         docker:
 9 |             - image: circleci/golang:1.12
10 |         steps:
11 |             - checkout
12 |             - restore_cache:
13 |                   keys:
14 |                       - go-mod-v4-{{ checksum "go.sum" }}
15 |             - run:
16 |                   name: Install deps
17 |                   command: |
18 |                         curl -sSfL https://raw.githubusercontent.com/golangci/golangci-lint/master/install.sh | sh -s -- -b $(go env GOPATH)/bin v1.23.3
19 |             - run:
20 |                   name: Run lints
21 |                   command: make lint
22 |             - run:
23 |                   name: Run unit tests
24 |                   command: make test
25 |             - codecov/upload:
26 |                 file: coverage.out
27 |             - save_cache:
28 |                   key: go-mod-v4-{{ checksum "go.sum" }}
29 |                   paths:
30 |                       - "/go/pkg/mod"
31 | 
32 | workflows:
33 |     version: 2.1
34 |     build:
35 |         jobs:
36 |             - build


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | .DS_Store
2 | .idea
3 | coverage.out
4 | coverage.html
5 | vendor
6 | examples/bench/server/server
7 | examples/bench/client/client
8 | 


--------------------------------------------------------------------------------
/.golangci.yml:
--------------------------------------------------------------------------------
 1 | run:
 2 |     deadline: 1m
 3 | 
 4 | linters:
 5 |     enable-all: true
 6 |     disable:
 7 |         - wsl
 8 | 
 9 | issues:
10 |     exclude:
11 |         - by other packages, and that stutters; consider calling this
12 | 


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
 1 | The MIT License (MIT)
 2 | 
 3 | Copyright (c) 2013-2014 Christopher Rosell
 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
13 | all 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
21 | THE SOFTWARE.
22 | 


--------------------------------------------------------------------------------
/Makefile:
--------------------------------------------------------------------------------
 1 | SHELL := /bin/sh
 2 | .DEFAULT_GOAL := default
 3 | MAKEFILE_PATH := $(abspath $(lastword $(MAKEFILE_LIST)))
 4 | CURRENT_DIR := $(patsubst %/,%,$(dir $(MAKEFILE_PATH)))
 5 | 
 6 | .PHONY: default
 7 | default: help
 8 | 
 9 | .PHONY: deps-install
10 | # @HELP Intsall project dependencies
11 | deps-install:
12 | 	@curl -sSfL https://raw.githubusercontent.com/golangci/golangci-lint/master/install.sh | sh -s -- -b $(shell go env GOPATH)/bin v1.23.3
13 | 	@go get -v github.com/golang/mock/mockgen
14 | 
15 | .PHONY: mocks-generate
16 | # @HELP Generate mocks for Go interfaces
17 | mocks-generate:
18 | 	@mockgen -source=hub.go -destination=mock/hub.go -package=mock
19 | 	@mockgen -source=client_factory.go -destination=mock/client_factory.go -package=mock
20 | 
21 | .PHONY: lint
22 | # @HELP Lint source code
23 | lint:
24 | 	@golangci-lint run --tests=false -v
25 | 
26 | .PHONY: test
27 | # Examples:
28 | # 	1. Run all tests:
29 | #		$ make test
30 | # 	2. Run all tests in a package:
31 | #		$ make test p=./path/to/pkg
32 | # 	3. Run a single test in a package:
33 | #		$ make test p=./path/to/pkg t=TestFunc/Name
34 | # @HELP Run tests
35 | test:
36 |     ifneq ($(p),)
37 |         ifneq ($(t),)
38 | 			@go test -cover -covermode=atomic -coverprofile=coverage.out -race ${p} -v -run ${t} && go tool cover -html=coverage.out -o coverage.html
39 |         else
40 | 			@go test -cover -covermode=atomic -coverprofile=coverage.out -race ${p} -v && go tool cover -html=coverage.out -o coverage.html
41 |         endif
42 |     else
43 | 		@go test -cover -covermode=atomic -coverprofile=coverage.out -race -v ./... && go tool cover -html=coverage.out -o coverage.html
44 |     endif
45 | 
46 | .PHONY: help
47 | # @HELP Print help
48 | help:
49 | 	@echo "\033[33mPlease use \`make <target>\` where <target> is one of:\033[0m"
50 | 	@sed -ne"/^# @HELP /{h;s/.*//;:d" -e"H;n;s/^# @HELP //;td" -e"s/:.*//;G;s/\\n# @HELP /---/;s/\\n/ /g;p;}" ${MAKEFILE_LIST}|LC_ALL='C' sort -f|awk -F --- -v n=$$(tput cols) -v i=32 -v a="$$(tput setaf 6)" -v z="$$(tput sgr0)" '{printf"%s%*s%s ",a,-i,$$1,z;m=split($$2,w," ");l=n-i;for(j=1;j<=m;j++){l-=length(w[j])+1;if(l<= 0){l=n-i-length(w[j])-1;printf"\n%*s ",-i," ";}printf"%s ",w[j];}printf"\n";}'
51 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
  1 | # WSPubSub
  2 | 
  3 | [![CircleCI Build Status](https://circleci.com/gh/kpeu3i/wspubsub.svg?style=shield)](https://circleci.com/gh/kpeu3i/wspubsub)
  4 | [![Coverage Status](https://codecov.io/gh/kpeu3i/wspubsub/graphs/badge.svg)](https://codecov.io/gh/kpeu3i/wspubsub/graphs/badge.svg)
  5 | [![GoDoc](https://godoc.org/github.com/kpeu3i/wspubsub?status.svg)](https://godoc.org/github.com/kpeu3i/wspubsub)
  6 | [![MIT Licensed](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
  7 | 
  8 | WSPubSub library is a Go implementation of channels based [pub/sub pattern](https://en.wikipedia.org/wiki/Publish%E2%80%93subscribe_pattern) over WebSocket protocol.
  9 | 
 10 | This library completely hides interaction with the transport layer like: connection upgrading, disconnecting, ping/pong etc.
 11 | Thanks to this, you can focus only on your tasks.
 12 | 
 13 | Client interaction with the library mainly occurs through the `hub` API. The only two steps to publish messages are required:
 14 | * Register receive handler using `hub.OnReceive`
 15 | * Subscribe clients using `hub.Subscribe`
 16 | 
 17 | Now you are ready to publish messages to different channels using `hub.Publish`.
 18 | Users who have been subscribed to those channels will receive the messages!
 19 | 
 20 | ## Install
 21 | 
 22 | Use go get to install the latest version of the library:
 23 | 
 24 |     go get github.com/kpeu3i/wspubsub
 25 | 
 26 | Next, include `wspubsub` in your application:
 27 | 
 28 | ```go
 29 | import "github.com/kpeu3i/wspubsub"
 30 | ```
 31 | 
 32 | ## Usage
 33 | 
 34 | A minimal working example is listed below:
 35 | 
 36 | ```go
 37 | package main
 38 | 
 39 | import (
 40 | 	"encoding/json"
 41 | 	"fmt"
 42 | 	"math/rand"
 43 | 	"os"
 44 | 	"os/signal"
 45 | 	"syscall"
 46 | 	"time"
 47 | 
 48 | 	"github.com/kpeu3i/wspubsub"
 49 | 	"github.com/pkg/errors"
 50 | )
 51 | 
 52 | type Message struct {
 53 | 	Command  string   `json:"command"`
 54 | 	Channels []string `json:"channels"`
 55 | }
 56 | 
 57 | func main() {
 58 | 	quit := make(chan os.Signal, 1)
 59 | 	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
 60 | 
 61 | 	channels := []string{"general", "public", "private"}
 62 | 	messageFormat := `{"now": %d}`
 63 | 
 64 | 	hub := wspubsub.NewDefaultHub()
 65 | 	defer hub.Close()
 66 | 
 67 | 	publishTicker := time.NewTicker(1 * time.Second)
 68 | 	defer publishTicker.Stop()
 69 | 
 70 | 	hub.OnReceive(func(clientID wspubsub.UUID, message wspubsub.Message) {
 71 | 		m := Message{}
 72 | 		err := json.Unmarshal(message.Payload, &m)
 73 | 		if err != nil {
 74 | 			hub.LogError(errors.Wrap(err, "hub.on_receive.unmarshal"))
 75 | 			return
 76 | 		}
 77 | 
 78 | 		switch m.Command {
 79 | 		case "SUBSCRIBE":
 80 | 			err := hub.Subscribe(clientID, m.Channels...)
 81 | 			if err != nil {
 82 | 				hub.LogError(errors.Wrap(err, "hub.on_receive.subscribe"))
 83 | 			}
 84 | 		case "UNSUBSCRIBE":
 85 | 			err := hub.Unsubscribe(clientID, m.Channels...)
 86 | 			if err != nil {
 87 | 				hub.LogError(errors.Wrap(err, "hub.on_receive.unsubscribe"))
 88 | 			}
 89 | 		}
 90 | 	})
 91 | 
 92 | 	go func() {
 93 | 		err := hub.ListenAndServe("localhost:8080", "/")
 94 | 		if err != nil {
 95 | 			hub.LogPanic(err)
 96 | 		}
 97 | 	}()
 98 | 
 99 | 	go func() {
100 | 		for range publishTicker.C {
101 | 			// Pick a random channel
102 | 			channel := channels[rand.Intn(len(channels))]
103 | 			message := wspubsub.NewTextMessageFromString(fmt.Sprintf(messageFormat, time.Now().Unix()))
104 | 			_, err := hub.Publish(message, channel)
105 | 			if err != nil {
106 | 				hub.LogPanic(err)
107 | 			}
108 | 
109 | 			hub.LogInfof("Published: channel=%s, message=%s\n", channel, string(message.Payload))
110 | 		}
111 | 	}()
112 | 
113 | 	<-quit
114 | }
115 | ````
116 | More examples you can find in [examples](https://github.com/kpeu3i/wspubsub/tree/master/examples/) directory.
117 | 
118 | ## Lint and test
119 | 
120 | First, you need to install project dependencies (golangci, mockgen, etc):
121 | 
122 |     $ make deps-install
123 | 
124 | Now, you are able to run:
125 | 
126 |     $ make lint
127 |     $ make test
128 | 
129 | If it required to update mocks use the following command:
130 | 
131 |     $ make mocks-generate
132 | 
133 | ## Contributing
134 | 
135 | Please don't hesitate to fork the project and send a pull request me to ask questions and share ideas.
136 | 
137 | 1. Fork it
138 | 2. Download the fork (`git clone https://github.com/kpeu3i/wspubsub.git && cd wspubsub`)
139 | 3. Create your feature branch (`git checkout -b my-new-feature`)
140 | 4. Make changes and add them (`git add .`)
141 | 5. Commit your changes (`git commit -m 'Add some feature'`)
142 | 6. Push to the branch (`git push origin my-new-feature`)
143 | 7. Create new pull request
144 | 
145 | ## License
146 | 
147 | This library is released under the MIT license. See the [LICENSE](https://github.com/kpeu3i/wspubsub/blob/master/LICENSE) file for details.
148 | 


--------------------------------------------------------------------------------
/client.go:
--------------------------------------------------------------------------------
  1 | package wspubsub
  2 | 
  3 | import (
  4 | 	"net/http"
  5 | 	"sync/atomic"
  6 | 	"time"
  7 | 
  8 | 	"github.com/pkg/errors"
  9 | )
 10 | 
 11 | // Client represents a connection to the WebSocket server.
 12 | type Client struct {
 13 | 	options        ClientOptions
 14 | 	id             UUID
 15 | 	upgrader       WebsocketConnectionUpgrader
 16 | 	logger         Logger
 17 | 	receiveHandler atomic.Value
 18 | 	errorHandler   atomic.Value
 19 | 	connection     WebsocketConnection
 20 | 	messages       chan Message
 21 | 	isConnected    bool
 22 | 	quit           chan struct{}
 23 | }
 24 | 
 25 | // ID returns unique client id.
 26 | func (c *Client) ID() UUID {
 27 | 	return c.id
 28 | }
 29 | 
 30 | // Connect upgrades the HTTP server connection to the WebSocket protocol.
 31 | func (c *Client) Connect(response http.ResponseWriter, request *http.Request) error {
 32 | 	if c.options.IsDebug {
 33 | 		now := time.Now()
 34 | 		defer func() {
 35 | 			end := time.Since(now)
 36 | 			if end > c.options.DebugFuncTimeLimit {
 37 | 				c.logger.Warnf("wspubsub.client.connect: took=%s", end)
 38 | 			}
 39 | 		}()
 40 | 	}
 41 | 
 42 | 	if c.isConnected {
 43 | 		return NewClientRepeatConnectError(c.id)
 44 | 	}
 45 | 
 46 | 	connection, err := c.upgrader.Upgrade(response, request)
 47 | 	if err != nil {
 48 | 		return errors.WithStack(NewClientConnectError(c.id, err))
 49 | 	}
 50 | 
 51 | 	c.connection = connection
 52 | 	c.isConnected = true
 53 | 
 54 | 	go c.runReader()
 55 | 	go c.runWriter()
 56 | 
 57 | 	return nil
 58 | }
 59 | 
 60 | // OnReceive registers a handler for incoming messages.
 61 | func (c *Client) OnReceive(handler ReceiveHandler) {
 62 | 	c.receiveHandler.Store(handler)
 63 | }
 64 | 
 65 | // OnError registers a handler for errors occurred while reading or writing connection.
 66 | func (c *Client) OnError(handler ErrorHandler) {
 67 | 	c.errorHandler.Store(handler)
 68 | }
 69 | 
 70 | // Send writes a message to client connection asynchronously.
 71 | func (c *Client) Send(message Message) error {
 72 | 	if c.options.IsDebug {
 73 | 		now := time.Now()
 74 | 		defer func() {
 75 | 			end := time.Since(now)
 76 | 			if end > c.options.DebugFuncTimeLimit {
 77 | 				c.logger.Warnf("wspubsub.client.send: took=%s", end)
 78 | 			}
 79 | 		}()
 80 | 	}
 81 | 
 82 | 	select {
 83 | 	case c.messages <- message:
 84 | 	default:
 85 | 		return errors.WithStack(NewClientSendBufferOverflowError(c.id))
 86 | 	}
 87 | 
 88 | 	return nil
 89 | }
 90 | 
 91 | // Close closes a client connection.
 92 | func (c *Client) Close() error {
 93 | 	if c.options.IsDebug {
 94 | 		now := time.Now()
 95 | 		defer func() {
 96 | 			end := time.Since(now)
 97 | 			if end > c.options.DebugFuncTimeLimit {
 98 | 				c.logger.Warnf("wspubsub.client.close: took=%s", end)
 99 | 			}
100 | 		}()
101 | 	}
102 | 
103 | 	if !c.isConnected {
104 | 		return nil
105 | 	}
106 | 
107 | 	defer func() {
108 | 		c.isConnected = false
109 | 	}()
110 | 
111 | 	c.quit <- struct{}{}
112 | 
113 | 	err := c.connection.Close()
114 | 	if err != nil {
115 | 		return errors.WithStack(err)
116 | 	}
117 | 
118 | 	return nil
119 | }
120 | 
121 | func (c *Client) runReader() {
122 | 	receiveHandler := c.receiveHandler.Load().(ReceiveHandler)
123 | 	errorHandler := c.errorHandler.Load().(ErrorHandler)
124 | 	for {
125 | 		message, err := c.connection.Read()
126 | 		if err != nil {
127 | 			err := errors.WithStack(NewClientReceiveError(c.id, message, err))
128 | 			errorHandler(c.id, err)
129 | 
130 | 			return
131 | 		}
132 | 
133 | 		receiveHandler(c.id, message)
134 | 	}
135 | }
136 | 
137 | func (c *Client) runWriter() {
138 | 	pingMessage := NewPingMessage()
139 | 	pingTicker := time.NewTicker(c.options.PingInterval)
140 | 	defer pingTicker.Stop()
141 | 
142 | 	pings := pingTicker.C
143 | 	messages := c.messages
144 | 	errorHandler := c.errorHandler.Load().(ErrorHandler)
145 | 	for {
146 | 		select {
147 | 		case <-c.quit:
148 | 			return
149 | 		case <-pings:
150 | 			err := c.connection.Write(pingMessage)
151 | 			if err != nil {
152 | 				err := errors.WithStack(NewClientPingError(c.id, pingMessage, err))
153 | 				errorHandler(c.id, err)
154 | 				pings = nil
155 | 			}
156 | 		case message := <-messages:
157 | 			err := c.connection.Write(message)
158 | 			if err != nil {
159 | 				err := errors.WithStack(NewClientSendError(c.id, message, err))
160 | 				errorHandler(c.id, err)
161 | 				messages = nil
162 | 			}
163 | 		}
164 | 	}
165 | }
166 | 
167 | // NewClient initializes a new Client.
168 | func NewClient(options ClientOptions, id UUID, upgrader WebsocketConnectionUpgrader, logger Logger) *Client {
169 | 	client := &Client{
170 | 		options:  options,
171 | 		id:       id,
172 | 		upgrader: upgrader,
173 | 		logger:   logger,
174 | 		messages: make(chan Message, options.SendBufferSize),
175 | 		quit:     make(chan struct{}),
176 | 	}
177 | 
178 | 	client.receiveHandler.Store(defaultReceiveHandler)
179 | 	client.errorHandler.Store(defaultErrorHandler)
180 | 
181 | 	return client
182 | }
183 | 


--------------------------------------------------------------------------------
/client_connect_error.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import (
 4 | 	"fmt"
 5 | 
 6 | 	"github.com/pkg/errors"
 7 | )
 8 | 
 9 | // ClientConnectError returned when HTTP connection can't be upgraded to WebSocket connection.
10 | type ClientConnectError struct {
11 | 	ID  UUID
12 | 	Err error
13 | }
14 | 
15 | // ClientConnectError implements an error interface.
16 | func (e *ClientConnectError) Error() string {
17 | 	return fmt.Sprintf("wspubsub: client failed to connect: id=%s, err=%s", e.ID, e.Err)
18 | }
19 | 
20 | // NewClientConnectError initializes a new ClientConnectError.
21 | func NewClientConnectError(id UUID, err error) *ClientConnectError {
22 | 	return &ClientConnectError{ID: id, Err: err}
23 | }
24 | 
25 | // IsClientConnectError checks if error type is ClientConnectError.
26 | func IsClientConnectError(err error) (*ClientConnectError, bool) {
27 | 	v, ok := errors.Cause(err).(*ClientConnectError)
28 | 
29 | 	return v, ok
30 | }
31 | 


--------------------------------------------------------------------------------
/client_connect_error_test.go:
--------------------------------------------------------------------------------
 1 | package wspubsub_test
 2 | 
 3 | import (
 4 | 	"errors"
 5 | 	"testing"
 6 | 
 7 | 	"github.com/kpeu3i/wspubsub"
 8 | 	"github.com/stretchr/testify/require"
 9 | )
10 | 
11 | func TestClientConnectError(t *testing.T) {
12 | 	rawErr := errors.New("TEST")
13 | 	err := wspubsub.NewClientConnectError(clientID, rawErr)
14 | 	require.Equal(t, clientID, err.ID)
15 | 	require.Equal(t, rawErr, err.Err)
16 | 	require.NotEmpty(t, clientID, err.Error())
17 | 
18 | 	e, ok := wspubsub.IsClientConnectError(err)
19 | 	require.NotNil(t, e)
20 | 	require.True(t, ok)
21 | 
22 | 	e, ok = wspubsub.IsClientConnectError(rawErr)
23 | 	require.Nil(t, e)
24 | 	require.False(t, ok)
25 | }
26 | 


--------------------------------------------------------------------------------
/client_factory.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import (
 4 | 	"net/http"
 5 | )
 6 | 
 7 | // WebsocketConnectionUpgrader upgrades HTTP connection to the WebSocket connection.
 8 | type WebsocketConnectionUpgrader interface {
 9 | 	Upgrade(w http.ResponseWriter, r *http.Request) (WebsocketConnection, error)
10 | }
11 | 
12 | // WebsocketConnection represents a WebSocket connection.
13 | type WebsocketConnection interface {
14 | 	Read() (Message, error)
15 | 	Write(message Message) error
16 | 	Close() error
17 | }
18 | 
19 | // UUIDGenerator generates UUID v4.
20 | type UUIDGenerator interface {
21 | 	GenerateV4() UUID
22 | }
23 | 
24 | // ClientFactory is responsible for creating a client.
25 | type ClientFactory struct {
26 | 	options       ClientOptions
27 | 	uuidGenerator UUIDGenerator
28 | 	upgrader      WebsocketConnectionUpgrader
29 | 	logger        Logger
30 | }
31 | 
32 | // Create returns a new client.
33 | func (f *ClientFactory) Create() WebsocketClient {
34 | 	return NewClient(f.options, f.uuidGenerator.GenerateV4(), f.upgrader, f.logger)
35 | }
36 | 
37 | // NewClientFactory initializes a new ClientFactory.
38 | func NewClientFactory(
39 | 	options ClientOptions,
40 | 	uuidGenerator UUIDGenerator,
41 | 	upgrader WebsocketConnectionUpgrader,
42 | 	logger Logger,
43 | ) *ClientFactory {
44 | 	factory := &ClientFactory{
45 | 		options:       options,
46 | 		uuidGenerator: uuidGenerator,
47 | 		upgrader:      upgrader,
48 | 		logger:        logger,
49 | 	}
50 | 
51 | 	return factory
52 | }
53 | 


--------------------------------------------------------------------------------
/client_factory_test.go:
--------------------------------------------------------------------------------
 1 | package wspubsub_test
 2 | 
 3 | import (
 4 | 	"testing"
 5 | 
 6 | 	"github.com/golang/mock/gomock"
 7 | 	"github.com/kpeu3i/wspubsub"
 8 | 	"github.com/kpeu3i/wspubsub/mock"
 9 | 	"github.com/stretchr/testify/require"
10 | )
11 | 
12 | func TestClientFactory_Create(t *testing.T) {
13 | 	ctrl := gomock.NewController(t)
14 | 	defer ctrl.Finish()
15 | 
16 | 	uuidGenerator := mock.NewMockUUIDGenerator(ctrl)
17 | 	uuidGenerator.
18 | 		EXPECT().
19 | 		GenerateV4().
20 | 		Return(clientID)
21 | 
22 | 	upgrader := mock.NewMockWebsocketConnectionUpgrader(ctrl)
23 | 	clientOptions := wspubsub.NewClientOptions()
24 | 	logger := mock.NewMockLogger(ctrl)
25 | 	factory := wspubsub.NewClientFactory(clientOptions, uuidGenerator, upgrader, logger)
26 | 	client := factory.Create()
27 | 
28 | 	require.NotNil(t, client)
29 | 	require.NotNil(t, client.ID())
30 | }
31 | 


--------------------------------------------------------------------------------
/client_not_found_error.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import (
 4 | 	"fmt"
 5 | 
 6 | 	"github.com/pkg/errors"
 7 | )
 8 | 
 9 | // ClientNotFoundError returned when client is not present in a storage.
10 | type ClientNotFoundError struct {
11 | 	ID UUID
12 | }
13 | 
14 | // ClientNotFoundError implements an error interface.
15 | func (e *ClientNotFoundError) Error() string {
16 | 	return fmt.Sprintf("wspubsub: client not found: id=%s", e.ID)
17 | }
18 | 
19 | // NewClientNotFoundError initializes a new ClientNotFoundError.
20 | func NewClientNotFoundError(id UUID) *ClientNotFoundError {
21 | 	return &ClientNotFoundError{ID: id}
22 | }
23 | 
24 | // IsClientNotFoundError checks if error type is ClientNotFoundError.
25 | func IsClientNotFoundError(err error) (*ClientNotFoundError, bool) {
26 | 	v, ok := errors.Cause(err).(*ClientNotFoundError)
27 | 
28 | 	return v, ok
29 | }
30 | 


--------------------------------------------------------------------------------
/client_not_found_error_test.go:
--------------------------------------------------------------------------------
 1 | package wspubsub_test
 2 | 
 3 | import (
 4 | 	"errors"
 5 | 	"testing"
 6 | 
 7 | 	"github.com/kpeu3i/wspubsub"
 8 | 	"github.com/stretchr/testify/require"
 9 | )
10 | 
11 | func TestClientNotFoundError(t *testing.T) {
12 | 	rawErr := errors.New("TEST")
13 | 	err := wspubsub.NewClientNotFoundError(clientID)
14 | 	require.Equal(t, clientID, err.ID)
15 | 	require.NotEmpty(t, err.Error())
16 | 
17 | 	e, ok := wspubsub.IsClientNotFoundError(err)
18 | 	require.NotNil(t, e)
19 | 	require.True(t, ok)
20 | 
21 | 	e, ok = wspubsub.IsClientNotFoundError(rawErr)
22 | 	require.Nil(t, e)
23 | 	require.False(t, ok)
24 | }
25 | 


--------------------------------------------------------------------------------
/client_options.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import "time"
 4 | 
 5 | // ClientOptions represents configuration of the client.
 6 | type ClientOptions struct {
 7 | 	// How often pings will be sent by the client.
 8 | 	PingInterval time.Duration
 9 | 
10 | 	// Max size of the buffer for messages which client should
11 | 	// write to a WebSocket connection.
12 | 	// Exceeding this size will cause an error.
13 | 	SendBufferSize int
14 | 
15 | 	// Enable/disable debug mode.
16 | 	IsDebug bool
17 | 
18 | 	// Function execution time limit in debug mode.
19 | 	// Exceeding this time limit will cause a new warn log message.
20 | 	DebugFuncTimeLimit time.Duration
21 | }
22 | 
23 | // NewClientOptions initializes a new ClientOptions.
24 | // nolint: gomnd
25 | func NewClientOptions() ClientOptions {
26 | 	return ClientOptions{
27 | 		PingInterval:       10 * time.Second,
28 | 		SendBufferSize:     1000,
29 | 		IsDebug:            false,
30 | 		DebugFuncTimeLimit: 1 * time.Millisecond,
31 | 	}
32 | }
33 | 


--------------------------------------------------------------------------------
/client_options_test.go:
--------------------------------------------------------------------------------
 1 | package wspubsub_test
 2 | 
 3 | import (
 4 | 	"testing"
 5 | 
 6 | 	"github.com/kpeu3i/wspubsub"
 7 | 	"github.com/stretchr/testify/require"
 8 | )
 9 | 
10 | func TestNewClientOptions(t *testing.T) {
11 | 	options := wspubsub.NewClientOptions()
12 | 	require.NotZero(t, options.PingInterval)
13 | 	require.NotZero(t, options.SendBufferSize)
14 | 	require.False(t, options.IsDebug)
15 | 	require.NotZero(t, options.DebugFuncTimeLimit)
16 | }
17 | 


--------------------------------------------------------------------------------
/client_ping_error.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import (
 4 | 	"fmt"
 5 | 
 6 | 	"github.com/pkg/errors"
 7 | )
 8 | 
 9 | // ClientPingError returned when ping message can't be written to a WebSocket connection.
10 | type ClientPingError struct {
11 | 	ID      UUID
12 | 	Message Message
13 | 	Err     error
14 | }
15 | 
16 | // ClientPingError implements an error interface.
17 | func (e *ClientPingError) Error() string {
18 | 	return fmt.Sprintf("wspubsub: client failed to send a ping message: id=%s, err=%s", e.ID, e.Err)
19 | }
20 | 
21 | // NewClientPingError initializes a new ClientPingError.
22 | func NewClientPingError(id UUID, message Message, err error) *ClientPingError {
23 | 	return &ClientPingError{ID: id, Message: message, Err: err}
24 | }
25 | 
26 | // IsClientPingError checks if error type is ClientPingError.
27 | func IsClientPingError(err error) (*ClientPingError, bool) {
28 | 	v, ok := errors.Cause(err).(*ClientPingError)
29 | 
30 | 	return v, ok
31 | }
32 | 


--------------------------------------------------------------------------------
/client_ping_error_test.go:
--------------------------------------------------------------------------------
 1 | package wspubsub_test
 2 | 
 3 | import (
 4 | 	"errors"
 5 | 	"testing"
 6 | 
 7 | 	"github.com/kpeu3i/wspubsub"
 8 | 	"github.com/stretchr/testify/require"
 9 | )
10 | 
11 | func TestClientPingError(t *testing.T) {
12 | 	message := wspubsub.NewTextMessageFromString("TEST")
13 | 	rawErr := errors.New("TEST")
14 | 	err := wspubsub.NewClientPingError(clientID, message, rawErr)
15 | 	require.Equal(t, clientID, err.ID)
16 | 	require.Equal(t, message, err.Message)
17 | 	require.Equal(t, rawErr, err.Err)
18 | 	require.NotEmpty(t, clientID, err.Error())
19 | 
20 | 	e, ok := wspubsub.IsClientPingError(err)
21 | 	require.NotNil(t, e)
22 | 	require.True(t, ok)
23 | 
24 | 	e, ok = wspubsub.IsClientPingError(rawErr)
25 | 	require.Nil(t, e)
26 | 	require.False(t, ok)
27 | }
28 | 


--------------------------------------------------------------------------------
/client_receive_error.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import (
 4 | 	"fmt"
 5 | 
 6 | 	"github.com/pkg/errors"
 7 | )
 8 | 
 9 | // ClientReceiveError returned when message can't be read from a WebSocket connection.
10 | type ClientReceiveError struct {
11 | 	ID      UUID
12 | 	Message Message
13 | 	Err     error
14 | }
15 | 
16 | // ClientReceiveError implements an error interface.
17 | func (e *ClientReceiveError) Error() string {
18 | 	return fmt.Sprintf("wspubsub: client failed to receieve a message: id=%s, err=%s", e.ID, e.Err)
19 | }
20 | 
21 | // NewClientReceiveError initializes a new ClientReceiveError.
22 | func NewClientReceiveError(id UUID, message Message, err error) *ClientReceiveError {
23 | 	return &ClientReceiveError{ID: id, Message: message, Err: err}
24 | }
25 | 
26 | // IsClientReceiveError checks if error type is ClientReceiveError.
27 | func IsClientReceiveError(err error) (*ClientReceiveError, bool) {
28 | 	v, ok := errors.Cause(err).(*ClientReceiveError)
29 | 
30 | 	return v, ok
31 | }
32 | 


--------------------------------------------------------------------------------
/client_receive_error_test.go:
--------------------------------------------------------------------------------
 1 | package wspubsub_test
 2 | 
 3 | import (
 4 | 	"errors"
 5 | 	"testing"
 6 | 
 7 | 	"github.com/kpeu3i/wspubsub"
 8 | 	"github.com/stretchr/testify/require"
 9 | )
10 | 
11 | func TestClientReceiveError(t *testing.T) {
12 | 	message := wspubsub.NewTextMessageFromString("TEST")
13 | 	rawErr := errors.New("TEST")
14 | 	err := wspubsub.NewClientReceiveError(clientID, message, rawErr)
15 | 	require.Equal(t, clientID, err.ID)
16 | 	require.Equal(t, message, err.Message)
17 | 	require.Equal(t, rawErr, err.Err)
18 | 	require.NotEmpty(t, clientID, err.Error())
19 | 
20 | 	e, ok := wspubsub.IsClientReceiveError(err)
21 | 	require.NotNil(t, e)
22 | 	require.True(t, ok)
23 | 
24 | 	e, ok = wspubsub.IsClientReceiveError(rawErr)
25 | 	require.Nil(t, e)
26 | 	require.False(t, ok)
27 | }
28 | 


--------------------------------------------------------------------------------
/client_repeat_connect_error.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import (
 4 | 	"fmt"
 5 | 
 6 | 	"github.com/pkg/errors"
 7 | )
 8 | 
 9 | // ClientRepeatConnectError returned when trying to connect an already connected client.
10 | type ClientRepeatConnectError struct {
11 | 	ID UUID
12 | }
13 | 
14 | // ClientRepeatConnectError implements an error interface.
15 | func (e *ClientRepeatConnectError) Error() string {
16 | 	return fmt.Sprintf("wspubsub: client is already connected: id=%s", e.ID)
17 | }
18 | 
19 | // NewClientRepeatConnectError initializes a new ClientRepeatConnectError.
20 | func NewClientRepeatConnectError(id UUID) *ClientRepeatConnectError {
21 | 	return &ClientRepeatConnectError{ID: id}
22 | }
23 | 
24 | // IsClientRepeatConnectError checks if error type is ClientRepeatConnectError.
25 | func IsClientRepeatConnectError(err error) (*ClientRepeatConnectError, bool) {
26 | 	v, ok := errors.Cause(err).(*ClientRepeatConnectError)
27 | 
28 | 	return v, ok
29 | }
30 | 


--------------------------------------------------------------------------------
/client_repeat_connect_error_test.go:
--------------------------------------------------------------------------------
 1 | package wspubsub_test
 2 | 
 3 | import (
 4 | 	"errors"
 5 | 	"testing"
 6 | 
 7 | 	"github.com/kpeu3i/wspubsub"
 8 | 	"github.com/stretchr/testify/require"
 9 | )
10 | 
11 | func TestClientRepeatConnectError(t *testing.T) {
12 | 	rawErr := errors.New("TEST")
13 | 	err := wspubsub.NewClientRepeatConnectError(clientID)
14 | 	require.Equal(t, clientID, err.ID)
15 | 	require.NotEmpty(t, clientID, err.Error())
16 | 
17 | 	e, ok := wspubsub.IsClientRepeatConnectError(err)
18 | 	require.NotNil(t, e)
19 | 	require.True(t, ok)
20 | 
21 | 	e, ok = wspubsub.IsClientRepeatConnectError(rawErr)
22 | 	require.Nil(t, e)
23 | 	require.False(t, ok)
24 | }
25 | 


--------------------------------------------------------------------------------
/client_send_buffer_overflow_error.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import (
 4 | 	"fmt"
 5 | 
 6 | 	"github.com/pkg/errors"
 7 | )
 8 | 
 9 | // ClientSendBufferOverflowError returned when client send buffer is full.
10 | type ClientSendBufferOverflowError struct {
11 | 	ID UUID
12 | }
13 | 
14 | // ClientSendBufferOverflowError implements an error interface.
15 | func (e *ClientSendBufferOverflowError) Error() string {
16 | 	return fmt.Sprintf("wspubsub: client send buffer is full: id=%s", e.ID)
17 | }
18 | 
19 | // NewClientSendBufferOverflowError initializes a new ClientSendBufferOverflowError.
20 | func NewClientSendBufferOverflowError(id UUID) *ClientSendBufferOverflowError {
21 | 	return &ClientSendBufferOverflowError{ID: id}
22 | }
23 | 
24 | // IsClientSendBufferOverflowError checks if error type is ClientSendBufferOverflowError.
25 | func IsClientSendBufferOverflowError(err error) (*ClientSendBufferOverflowError, bool) {
26 | 	v, ok := errors.Cause(err).(*ClientSendBufferOverflowError)
27 | 
28 | 	return v, ok
29 | }
30 | 


--------------------------------------------------------------------------------
/client_send_buffer_overflow_error_test.go:
--------------------------------------------------------------------------------
 1 | package wspubsub_test
 2 | 
 3 | import (
 4 | 	"errors"
 5 | 	"testing"
 6 | 
 7 | 	"github.com/kpeu3i/wspubsub"
 8 | 	"github.com/stretchr/testify/require"
 9 | )
10 | 
11 | func TestClientSendBufferOverflowError(t *testing.T) {
12 | 	rawErr := errors.New("TEST")
13 | 	err := wspubsub.NewClientSendBufferOverflowError(clientID)
14 | 	require.Equal(t, clientID, err.ID)
15 | 	require.NotEmpty(t, clientID, err.Error())
16 | 
17 | 	e, ok := wspubsub.IsClientSendBufferOverflowError(err)
18 | 	require.NotNil(t, e)
19 | 	require.True(t, ok)
20 | 
21 | 	e, ok = wspubsub.IsClientSendBufferOverflowError(rawErr)
22 | 	require.Nil(t, e)
23 | 	require.False(t, ok)
24 | }
25 | 


--------------------------------------------------------------------------------
/client_send_error.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import (
 4 | 	"fmt"
 5 | 
 6 | 	"github.com/pkg/errors"
 7 | )
 8 | 
 9 | // ClientSendError returned when a message (text or binary) can't be written to a WebSocket connection.
10 | type ClientSendError struct {
11 | 	ID      UUID
12 | 	Message Message
13 | 	Err     error
14 | }
15 | 
16 | // ClientSendError implements an error interface.
17 | func (e *ClientSendError) Error() string {
18 | 	return fmt.Sprintf("wspubsub: client failed to send a message: id=%s, err=%s", e.ID, e.Err)
19 | }
20 | 
21 | // NewClientSendError initializes a new ClientSendError.
22 | func NewClientSendError(id UUID, message Message, err error) *ClientSendError {
23 | 	return &ClientSendError{ID: id, Message: message, Err: err}
24 | }
25 | 
26 | // IsClientSendError checks if error type is ClientSendError.
27 | func IsClientSendError(err error) (*ClientSendError, bool) {
28 | 	v, ok := errors.Cause(err).(*ClientSendError)
29 | 
30 | 	return v, ok
31 | }
32 | 


--------------------------------------------------------------------------------
/client_send_error_test.go:
--------------------------------------------------------------------------------
 1 | package wspubsub_test
 2 | 
 3 | import (
 4 | 	"errors"
 5 | 	"testing"
 6 | 
 7 | 	"github.com/kpeu3i/wspubsub"
 8 | 	"github.com/stretchr/testify/require"
 9 | )
10 | 
11 | func TestClientSendError(t *testing.T) {
12 | 	message := wspubsub.NewTextMessageFromString("TEST")
13 | 	rawErr := errors.New("TEST")
14 | 	err := wspubsub.NewClientSendError(clientID, message, rawErr)
15 | 	require.Equal(t, clientID, err.ID)
16 | 	require.Equal(t, message, err.Message)
17 | 	require.Equal(t, rawErr, err.Err)
18 | 	require.NotEmpty(t, clientID, err.Error())
19 | 
20 | 	e, ok := wspubsub.IsClientSendError(err)
21 | 	require.NotNil(t, e)
22 | 	require.True(t, ok)
23 | 
24 | 	e, ok = wspubsub.IsClientSendError(rawErr)
25 | 	require.Nil(t, e)
26 | 	require.False(t, ok)
27 | }
28 | 


--------------------------------------------------------------------------------
/client_store.go:
--------------------------------------------------------------------------------
  1 | package wspubsub
  2 | 
  3 | import (
  4 | 	"sync"
  5 | 	"time"
  6 | 
  7 | 	"github.com/cespare/xxhash/v2"
  8 | 	"github.com/pkg/errors"
  9 | )
 10 | 
 11 | // IterateFunc is the type of the function called for
 12 | // each client visited by Find.
 13 | type IterateFunc func(client WebsocketClient) error
 14 | 
 15 | type clientsBuffer struct {
 16 | 	clients []WebsocketClient
 17 | }
 18 | 
 19 | // ClientStore represents the storage of clients.
 20 | type ClientStore struct {
 21 | 	options           ClientStoreOptions
 22 | 	logger            Logger
 23 | 	clientsShardList  []*clientStoreClientsShard
 24 | 	channelsShardList []*clientStoreChannelsShard
 25 | 	clientsPool       sync.Pool
 26 | }
 27 | 
 28 | // Get returns client by its ID.
 29 | func (s *ClientStore) Get(clientID UUID) (WebsocketClient, error) {
 30 | 	if s.options.IsDebug {
 31 | 		now := time.Now()
 32 | 		defer func() {
 33 | 			end := time.Since(now)
 34 | 			if end > s.options.DebugFuncTimeLimit {
 35 | 				s.logger.Warnf("wspubsub.client_store.get: took=%s", end)
 36 | 			}
 37 | 		}()
 38 | 	}
 39 | 
 40 | 	clientsShard := s.clientsShard(clientID)
 41 | 
 42 | 	client, err := clientsShard.Get(clientID)
 43 | 	if err != nil {
 44 | 		return nil, errors.WithStack(err)
 45 | 	}
 46 | 
 47 | 	return client, nil
 48 | }
 49 | 
 50 | // Set puts client to storage.
 51 | func (s *ClientStore) Set(client WebsocketClient) {
 52 | 	if s.options.IsDebug {
 53 | 		now := time.Now()
 54 | 		defer func() {
 55 | 			end := time.Since(now)
 56 | 			if end > s.options.DebugFuncTimeLimit {
 57 | 				s.logger.Warnf("wspubsub.client_store.set: took=%s", end)
 58 | 			}
 59 | 		}()
 60 | 	}
 61 | 
 62 | 	clientsShard := s.clientsShard(client.ID())
 63 | 	clientsShard.Set(client)
 64 | }
 65 | 
 66 | // Unset removes client from storage by its ID
 67 | func (s *ClientStore) Unset(clientID UUID) error {
 68 | 	if s.options.IsDebug {
 69 | 		now := time.Now()
 70 | 		defer func() {
 71 | 			end := time.Since(now)
 72 | 			if end > s.options.DebugFuncTimeLimit {
 73 | 				s.logger.Warnf("wspubsub.client_store.unset: took=%s", end)
 74 | 			}
 75 | 		}()
 76 | 	}
 77 | 
 78 | 	clientsShard := s.clientsShard(clientID)
 79 | 	clientsShard.Unset(clientID)
 80 | 
 81 | 	for _, channelsShard := range s.channelsShardList {
 82 | 		channelsShard.Unlink(clientID)
 83 | 	}
 84 | 
 85 | 	return nil
 86 | }
 87 | 
 88 | // Count returns the total number of clients in specified channel(-s).
 89 | func (s *ClientStore) Count(channels ...string) int {
 90 | 	if s.options.IsDebug {
 91 | 		now := time.Now()
 92 | 		defer func() {
 93 | 			end := time.Since(now)
 94 | 			if end > s.options.DebugFuncTimeLimit {
 95 | 				s.logger.Warnf("wspubsub.client_store.count: took=%s", end)
 96 | 			}
 97 | 		}()
 98 | 	}
 99 | 
100 | 	count := 0
101 | 	if len(channels) == 0 {
102 | 		for _, clientsShard := range s.clientsShardList {
103 | 			count += clientsShard.Count()
104 | 		}
105 | 
106 | 		return count
107 | 	}
108 | 
109 | 	for _, channel := range channels {
110 | 		channelsShard := s.channelsShard(channel)
111 | 		count += channelsShard.Count(channel)
112 | 	}
113 | 
114 | 	return count
115 | }
116 | 
117 | // Find iterates over clients who subscribed on specified channel(-s).
118 | func (s *ClientStore) Find(fn IterateFunc, channels ...string) error {
119 | 	if s.options.IsDebug {
120 | 		now := time.Now()
121 | 		defer func() {
122 | 			end := time.Since(now)
123 | 			if end > s.options.DebugFuncTimeLimit {
124 | 				s.logger.Warnf("wspubsub.client_store.find: took=%s", end)
125 | 			}
126 | 		}()
127 | 	}
128 | 
129 | 	buff := s.clientsPool.Get().(*clientsBuffer)
130 | 
131 | 	if len(channels) == 0 {
132 | 		for _, clientsShard := range s.clientsShardList {
133 | 			clientsShard.Iterate(func(client WebsocketClient) {
134 | 				buff.clients = append(buff.clients, client)
135 | 			})
136 | 		}
137 | 	} else {
138 | 		for _, channel := range channels {
139 | 			channelsShard := s.channelsShard(channel)
140 | 			channelsShard.Iterate(channel, func(client WebsocketClient) {
141 | 				buff.clients = append(buff.clients, client)
142 | 			})
143 | 		}
144 | 	}
145 | 
146 | 	for _, client := range buff.clients {
147 | 		err := fn(client)
148 | 		if err != nil {
149 | 			return errors.WithStack(err)
150 | 		}
151 | 	}
152 | 
153 | 	buff.clients = buff.clients[:0]
154 | 	s.clientsPool.Put(buff)
155 | 
156 | 	return nil
157 | }
158 | 
159 | // CountChannels return a list of channels linked with the client.
160 | func (s *ClientStore) Channels(clientID UUID) ([]string, error) {
161 | 	if s.options.IsDebug {
162 | 		now := time.Now()
163 | 		defer func() {
164 | 			end := time.Since(now)
165 | 			if end > s.options.DebugFuncTimeLimit {
166 | 				s.logger.Warnf("wspubsub.client_store.channels: took=%s", end)
167 | 			}
168 | 		}()
169 | 	}
170 | 
171 | 	clientsShard := s.clientsShard(clientID)
172 | 
173 | 	_, channels, ok := clientsShard.Channels(clientID)
174 | 	if !ok {
175 | 		return nil, errors.WithStack(NewClientNotFoundError(clientID))
176 | 	}
177 | 
178 | 	return channels, nil
179 | }
180 | 
181 | // CountChannels return the total number of channels linked with the client.
182 | func (s *ClientStore) CountChannels(clientID UUID) (int, error) {
183 | 	if s.options.IsDebug {
184 | 		now := time.Now()
185 | 		defer func() {
186 | 			end := time.Since(now)
187 | 			if end > s.options.DebugFuncTimeLimit {
188 | 				s.logger.Warnf("wspubsub.client_store.count_channels: took=%s", end)
189 | 			}
190 | 		}()
191 | 	}
192 | 
193 | 	clientsShard := s.clientsShard(clientID)
194 | 
195 | 	_, count, ok := clientsShard.CountChannels(clientID)
196 | 	if !ok {
197 | 		return 0, errors.WithStack(NewClientNotFoundError(clientID))
198 | 	}
199 | 
200 | 	return count, nil
201 | }
202 | 
203 | // SetChannels links the client with specified channel(-s).
204 | func (s *ClientStore) SetChannels(clientID UUID, channels ...string) error {
205 | 	if s.options.IsDebug {
206 | 		now := time.Now()
207 | 		defer func() {
208 | 			end := time.Since(now)
209 | 			if end > s.options.DebugFuncTimeLimit {
210 | 				s.logger.Warnf("wspubsub.client_store.set_channels: took=%s", end)
211 | 			}
212 | 		}()
213 | 	}
214 | 
215 | 	if len(channels) == 0 {
216 | 		return nil
217 | 	}
218 | 
219 | 	clientsShard := s.clientsShard(clientID)
220 | 
221 | 	client, ok := clientsShard.SetChannels(clientID, channels...)
222 | 	if !ok {
223 | 		return errors.WithStack(NewClientNotFoundError(clientID))
224 | 	}
225 | 
226 | 	for _, channel := range channels {
227 | 		channelsShard := s.channelsShard(channel)
228 | 		channelsShard.Link(client, channel)
229 | 	}
230 | 
231 | 	return nil
232 | }
233 | 
234 | // SetChannels unlinks the client from specified channel(-s).
235 | // If channels were not specified then the client will be
236 | // unlinked from all channels.
237 | func (s *ClientStore) UnsetChannels(clientID UUID, channels ...string) error {
238 | 	if s.options.IsDebug {
239 | 		now := time.Now()
240 | 		defer func() {
241 | 			end := time.Since(now)
242 | 			if end > s.options.DebugFuncTimeLimit {
243 | 				s.logger.Warnf("wspubsub.client_store.unset_channels: took=%s", end)
244 | 			}
245 | 		}()
246 | 	}
247 | 
248 | 	clientsShard := s.clientsShard(clientID)
249 | 
250 | 	_, ok := clientsShard.UnsetChannels(clientID, channels...)
251 | 	if !ok {
252 | 		return errors.WithStack(NewClientNotFoundError(clientID))
253 | 	}
254 | 
255 | 	if len(channels) == 0 {
256 | 		for _, channelsShard := range s.channelsShardList {
257 | 			channelsShard.Unlink(clientID)
258 | 		}
259 | 
260 | 		return nil
261 | 	}
262 | 
263 | 	for _, channel := range channels {
264 | 		channelsShard := s.channelsShard(channel)
265 | 		channelsShard.Unlink(clientID, channel)
266 | 	}
267 | 
268 | 	return nil
269 | }
270 | 
271 | func (s *ClientStore) clientsShard(clientID UUID) *clientStoreClientsShard {
272 | 	index := xxhash.Sum64(clientID.Bytes()) % uint64(s.options.ClientShards.Count)
273 | 
274 | 	return s.clientsShardList[index]
275 | }
276 | 
277 | func (s *ClientStore) channelsShard(channel string) *clientStoreChannelsShard {
278 | 	index := xxhash.Sum64String(channel) % uint64(s.options.ChannelShards.Count)
279 | 
280 | 	return s.channelsShardList[index]
281 | }
282 | 
283 | // NewClientStore initializes a new ClientStore.
284 | func NewClientStore(options ClientStoreOptions, logger Logger) *ClientStore {
285 | 	clientList := &ClientStore{
286 | 		options:           options,
287 | 		logger:            logger,
288 | 		clientsShardList:  make([]*clientStoreClientsShard, options.ClientShards.Count),
289 | 		channelsShardList: make([]*clientStoreChannelsShard, options.ChannelShards.Count),
290 | 		clientsPool: sync.Pool{
291 | 			New: func() interface{} {
292 | 				return &clientsBuffer{clients: make([]WebsocketClient, 0, options.ClientShards.Size)}
293 | 			},
294 | 		},
295 | 	}
296 | 
297 | 	for i := 0; i < options.ClientShards.Count; i++ {
298 | 		clientList.clientsShardList[i] = newClientStoreClientsShard(
299 | 			options.ClientShards.Size,
300 | 			options.ClientShards.BucketSize,
301 | 		)
302 | 	}
303 | 
304 | 	for i := 0; i < options.ChannelShards.Count; i++ {
305 | 		clientList.channelsShardList[i] = newClientStoreChannelsShard(
306 | 			options.ChannelShards.Size,
307 | 			options.ChannelShards.BucketSize,
308 | 		)
309 | 	}
310 | 
311 | 	return clientList
312 | }
313 | 


--------------------------------------------------------------------------------
/client_store_channels_shard.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import (
 4 | 	"sync"
 5 | )
 6 | 
 7 | type clientStoreChannelsShardBucket map[UUID]WebsocketClient
 8 | 
 9 | type clientStoreChannelsShard struct {
10 | 	bucketSize int
11 | 	mu         sync.RWMutex
12 | 	clients    map[string]clientStoreChannelsShardBucket
13 | }
14 | 
15 | func (s *clientStoreChannelsShard) Link(client WebsocketClient, channel string) {
16 | 	s.mu.Lock()
17 | 	if _, ok := s.clients[channel]; !ok {
18 | 		s.clients[channel] = make(clientStoreChannelsShardBucket, s.bucketSize)
19 | 	}
20 | 	s.clients[channel][client.ID()] = client
21 | 	s.mu.Unlock()
22 | }
23 | 
24 | func (s *clientStoreChannelsShard) Unlink(clientID UUID, channels ...string) {
25 | 	if len(channels) == 0 {
26 | 		s.mu.Lock()
27 | 		for channel := range s.clients {
28 | 			delete(s.clients[channel], clientID)
29 | 		}
30 | 		s.mu.Unlock()
31 | 
32 | 		return
33 | 	}
34 | 
35 | 	s.mu.Lock()
36 | 	for _, channel := range channels {
37 | 		delete(s.clients[channel], clientID)
38 | 	}
39 | 	s.mu.Unlock()
40 | }
41 | 
42 | func (s *clientStoreChannelsShard) Count(channels ...string) int {
43 | 	count := 0
44 | 
45 | 	if len(channels) == 0 {
46 | 		s.mu.RLock()
47 | 		for channel := range s.clients {
48 | 			count += len(s.clients[channel])
49 | 		}
50 | 		s.mu.RUnlock()
51 | 
52 | 		return count
53 | 	}
54 | 
55 | 	s.mu.RLock()
56 | 	for _, channel := range channels {
57 | 		count += len(s.clients[channel])
58 | 	}
59 | 	s.mu.RUnlock()
60 | 
61 | 	return count
62 | }
63 | 
64 | func (s *clientStoreChannelsShard) Iterate(channel string, iterateFunc func(client WebsocketClient)) {
65 | 	s.mu.RLock()
66 | 	if clients, ok := s.clients[channel]; ok {
67 | 		for _, client := range clients {
68 | 			iterateFunc(client)
69 | 		}
70 | 	}
71 | 	s.mu.RUnlock()
72 | }
73 | 
74 | func newClientStoreChannelsShard(size int, bucketSize int) *clientStoreChannelsShard {
75 | 	shard := &clientStoreChannelsShard{
76 | 		bucketSize: bucketSize,
77 | 		clients:    make(map[string]clientStoreChannelsShardBucket, size),
78 | 	}
79 | 
80 | 	return shard
81 | }
82 | 


--------------------------------------------------------------------------------
/client_store_clients_shard.go:
--------------------------------------------------------------------------------
  1 | package wspubsub
  2 | 
  3 | import (
  4 | 	"sync"
  5 | 
  6 | 	"github.com/pkg/errors"
  7 | )
  8 | 
  9 | type clientStoreShardValue struct {
 10 | 	client   WebsocketClient
 11 | 	channels map[string]struct{}
 12 | }
 13 | 
 14 | type clientStoreClientsShard struct {
 15 | 	bucketSize int
 16 | 	mu         sync.RWMutex
 17 | 	values     map[UUID]clientStoreShardValue
 18 | }
 19 | 
 20 | func (s *clientStoreClientsShard) Get(clientID UUID) (WebsocketClient, error) {
 21 | 	s.mu.RLock()
 22 | 	value, ok := s.values[clientID]
 23 | 	s.mu.RUnlock()
 24 | 	if !ok {
 25 | 		return nil, errors.WithStack(NewClientNotFoundError(clientID))
 26 | 	}
 27 | 
 28 | 	return value.client, nil
 29 | }
 30 | 
 31 | func (s *clientStoreClientsShard) Set(client WebsocketClient) {
 32 | 	s.mu.Lock()
 33 | 	s.values[client.ID()] = clientStoreShardValue{client: client, channels: make(map[string]struct{}, s.bucketSize)}
 34 | 	s.mu.Unlock()
 35 | }
 36 | 
 37 | func (s *clientStoreClientsShard) Unset(clientID UUID) {
 38 | 	s.mu.Lock()
 39 | 	delete(s.values, clientID)
 40 | 	s.mu.Unlock()
 41 | }
 42 | 
 43 | func (s *clientStoreClientsShard) Count() int {
 44 | 	s.mu.RLock()
 45 | 	count := len(s.values)
 46 | 	s.mu.RUnlock()
 47 | 
 48 | 	return count
 49 | }
 50 | 
 51 | func (s *clientStoreClientsShard) Channels(clientID UUID) (WebsocketClient, []string, bool) {
 52 | 	var (
 53 | 		client   WebsocketClient
 54 | 		channels []string
 55 | 	)
 56 | 
 57 | 	s.mu.RLock()
 58 | 	value, ok := s.values[clientID]
 59 | 	if ok {
 60 | 		client = value.client
 61 | 		channels = make([]string, 0, len(value.channels))
 62 | 		for channel := range value.channels {
 63 | 			channels = append(channels, channel)
 64 | 		}
 65 | 	}
 66 | 	s.mu.RUnlock()
 67 | 
 68 | 	return client, channels, ok
 69 | }
 70 | 
 71 | func (s *clientStoreClientsShard) CountChannels(clientID UUID) (WebsocketClient, int, bool) {
 72 | 	var (
 73 | 		client WebsocketClient
 74 | 		count  int
 75 | 	)
 76 | 
 77 | 	s.mu.RLock()
 78 | 	value, ok := s.values[clientID]
 79 | 	if ok {
 80 | 		client = value.client
 81 | 		count = len(value.channels)
 82 | 	}
 83 | 	s.mu.RUnlock()
 84 | 
 85 | 	return client, count, ok
 86 | }
 87 | 
 88 | func (s *clientStoreClientsShard) SetChannels(clientID UUID, channels ...string) (WebsocketClient, bool) {
 89 | 	var client WebsocketClient
 90 | 
 91 | 	s.mu.Lock()
 92 | 	value, ok := s.values[clientID]
 93 | 	if ok {
 94 | 		client = value.client
 95 | 		for _, channel := range channels {
 96 | 			value.channels[channel] = struct{}{}
 97 | 		}
 98 | 	}
 99 | 	s.mu.Unlock()
100 | 
101 | 	return client, ok
102 | }
103 | 
104 | func (s *clientStoreClientsShard) UnsetChannels(clientID UUID, channels ...string) (WebsocketClient, bool) {
105 | 	var client WebsocketClient
106 | 
107 | 	s.mu.Lock()
108 | 	value, ok := s.values[clientID]
109 | 	if ok {
110 | 		client = value.client
111 | 		if len(channels) == 0 {
112 | 			for channel := range value.channels {
113 | 				delete(value.channels, channel)
114 | 			}
115 | 		} else {
116 | 			for _, channel := range channels {
117 | 				delete(value.channels, channel)
118 | 			}
119 | 		}
120 | 	}
121 | 	s.mu.Unlock()
122 | 
123 | 	return client, ok
124 | }
125 | 
126 | func (s *clientStoreClientsShard) Iterate(iterateFunc func(client WebsocketClient)) {
127 | 	s.mu.RLock()
128 | 	for _, value := range s.values {
129 | 		iterateFunc(value.client)
130 | 	}
131 | 	s.mu.RUnlock()
132 | }
133 | 
134 | func newClientStoreClientsShard(size int, bucketSize int) *clientStoreClientsShard {
135 | 	return &clientStoreClientsShard{bucketSize: bucketSize, values: make(map[UUID]clientStoreShardValue, size)}
136 | }
137 | 


--------------------------------------------------------------------------------
/client_store_options.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import "time"
 4 | 
 5 | // ClientStoreOptions represents configuration of the storage.
 6 | type ClientStoreOptions struct {
 7 | 	ClientShards struct {
 8 | 		// Total number of shards
 9 | 		Count int
10 | 
11 | 		// Size of shard
12 | 		Size int
13 | 
14 | 		// Size of a bucket in shard
15 | 		BucketSize int
16 | 	}
17 | 	ChannelShards struct {
18 | 		// Total number of shards
19 | 		Count int
20 | 
21 | 		// Size of shard
22 | 		Size int
23 | 
24 | 		// Size of a bucket in shard
25 | 		BucketSize int
26 | 	}
27 | 
28 | 	// Enable/disable debug mode.
29 | 	IsDebug bool
30 | 
31 | 	// Function execution time limit in debug mode.
32 | 	// Exceeding this time limit will cause a new warn log message.
33 | 	DebugFuncTimeLimit time.Duration
34 | }
35 | 
36 | // NewClientStoreOptions initializes a new ClientStoreOptions.
37 | // nolint: gomnd
38 | func NewClientStoreOptions() ClientStoreOptions {
39 | 	options := ClientStoreOptions{
40 | 		IsDebug:            false,
41 | 		DebugFuncTimeLimit: 1 * time.Millisecond,
42 | 	}
43 | 
44 | 	options.ClientShards.Count = 128
45 | 	options.ClientShards.Size = 10000
46 | 	options.ChannelShards.BucketSize = 100
47 | 
48 | 	options.ChannelShards.Count = 16
49 | 	options.ChannelShards.Size = 100
50 | 	options.ChannelShards.BucketSize = 10000
51 | 
52 | 	return options
53 | }
54 | 


--------------------------------------------------------------------------------
/client_store_options_test.go:
--------------------------------------------------------------------------------
 1 | package wspubsub_test
 2 | 
 3 | import (
 4 | 	"testing"
 5 | 
 6 | 	"github.com/kpeu3i/wspubsub"
 7 | 	"github.com/stretchr/testify/require"
 8 | )
 9 | 
10 | func TestNewClientStoreOptions(t *testing.T) {
11 | 	options := wspubsub.NewClientStoreOptions()
12 | 	require.NotZero(t, options.ClientShards.Count)
13 | 	require.NotZero(t, options.ClientShards.Size)
14 | 	require.NotZero(t, options.ChannelShards.Count)
15 | 	require.NotZero(t, options.ChannelShards.Size)
16 | 	require.NotZero(t, options.ChannelShards.BucketSize)
17 | 	require.False(t, options.IsDebug)
18 | 	require.NotZero(t, options.DebugFuncTimeLimit)
19 | }
20 | 


--------------------------------------------------------------------------------
/client_store_test.go:
--------------------------------------------------------------------------------
  1 | package wspubsub_test
  2 | 
  3 | import (
  4 | 	"math/rand"
  5 | 	"sort"
  6 | 	"testing"
  7 | 	"time"
  8 | 
  9 | 	"github.com/golang/mock/gomock"
 10 | 	"github.com/kpeu3i/wspubsub"
 11 | 	"github.com/kpeu3i/wspubsub/mock"
 12 | 	"github.com/pkg/errors"
 13 | 	"github.com/stretchr/testify/require"
 14 | )
 15 | 
 16 | func init() {
 17 | 	rand.Seed(time.Now().UnixNano())
 18 | }
 19 | 
 20 | func TestClientStore_Clients(t *testing.T) {
 21 | 	ctrl := gomock.NewController(t)
 22 | 	defer ctrl.Finish()
 23 | 
 24 | 	logger := mock.NewMockLogger(ctrl)
 25 | 	upgrader := mock.NewMockWebsocketConnectionUpgrader(ctrl)
 26 | 
 27 | 	clientsNum := 10
 28 | 	clients := make(map[wspubsub.UUID]*wspubsub.Client, clientsNum)
 29 | 	clientList := make([]*wspubsub.Client, 0, clientsNum)
 30 | 	for i := 0; i < clientsNum; i++ {
 31 | 		cid := wspubsub.UUID([16]byte{byte(i), 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16})
 32 | 		clientOptions := wspubsub.NewClientOptions()
 33 | 		client := wspubsub.NewClient(clientOptions, cid, upgrader, logger)
 34 | 		clients[cid] = client
 35 | 		clientList = append(clientList, client)
 36 | 	}
 37 | 
 38 | 	clientStoreOptions := wspubsub.NewClientStoreOptions()
 39 | 	clientStore := wspubsub.NewClientStore(clientStoreOptions, logger)
 40 | 
 41 | 	t.Run("Check default state of a storage", func(t *testing.T) {
 42 | 		client, err := clientStore.Get(clientID)
 43 | 		require.Nil(t, client)
 44 | 		require.Equal(t, wspubsub.NewClientNotFoundError(clientID), errors.Cause(err).(*wspubsub.ClientNotFoundError))
 45 | 		require.Equal(t, 0, clientStore.Count())
 46 | 
 47 | 		err = clientStore.Unset(clientID)
 48 | 		require.NoError(t, err)
 49 | 	})
 50 | 
 51 | 	t.Run("Fill a storage", func(t *testing.T) {
 52 | 		i := 1
 53 | 		for cid, client := range clients {
 54 | 			clientStore.Set(client)
 55 | 			storedClient, err := clientStore.Get(cid)
 56 | 			require.NoError(t, err)
 57 | 			require.Equal(t, client, storedClient)
 58 | 			require.Equal(t, i, clientStore.Count())
 59 | 
 60 | 			i++
 61 | 		}
 62 | 	})
 63 | 
 64 | 	t.Run("Remove a client from a storage", func(t *testing.T) {
 65 | 		removedClientID := clientList[0].ID()
 66 | 		err := clientStore.Unset(removedClientID)
 67 | 		require.NoError(t, err)
 68 | 		require.Equal(t, len(clientList)-1, clientStore.Count())
 69 | 
 70 | 		client, err := clientStore.Get(removedClientID)
 71 | 		require.Nil(t, client)
 72 | 		require.Equal(t, wspubsub.NewClientNotFoundError(removedClientID), errors.Cause(err).(*wspubsub.ClientNotFoundError))
 73 | 
 74 | 		for cid, client := range clients {
 75 | 			if client.ID() == removedClientID {
 76 | 				continue
 77 | 			}
 78 | 
 79 | 			storedClient, err := clientStore.Get(cid)
 80 | 			require.NoError(t, err)
 81 | 			require.Equal(t, client, storedClient)
 82 | 		}
 83 | 	})
 84 | }
 85 | 
 86 | func TestClientStore_Find(t *testing.T) {
 87 | 	ctrl := gomock.NewController(t)
 88 | 	defer ctrl.Finish()
 89 | 
 90 | 	logger := mock.NewMockLogger(ctrl)
 91 | 	upgrader := mock.NewMockWebsocketConnectionUpgrader(ctrl)
 92 | 
 93 | 	clientStoreOptions := wspubsub.NewClientStoreOptions()
 94 | 	clientStore := wspubsub.NewClientStore(clientStoreOptions, logger)
 95 | 
 96 | 	numClients := 100
 97 | 	numClientsInChannels := 0
 98 | 	availableChannels := []string{"X", "Y", "Z"}
 99 | 	clients := map[string]map[wspubsub.UUID]*wspubsub.Client{}
100 | 
101 | 	for i := 0; i < numClients; i++ {
102 | 		cid := wspubsub.UUID([16]byte{byte(i), 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16})
103 | 		clientOptions := wspubsub.NewClientOptions()
104 | 		client := wspubsub.NewClient(clientOptions, cid, upgrader, logger)
105 | 		clientStore.Set(client)
106 | 
107 | 		numChannels := rand.Intn(len(availableChannels)-1) + 1
108 | 		channels := make([]string, 0, numChannels)
109 | 		for j := 0; j < numChannels; j++ {
110 | 			channels = append(channels, availableChannels[rand.Intn(len(availableChannels))])
111 | 		}
112 | 
113 | 		err := clientStore.SetChannels(cid, channels...)
114 | 		require.NoError(t, err)
115 | 
116 | 		for _, channel := range channels {
117 | 			if _, ok := clients[channel]; !ok {
118 | 				clients[channel] = make(map[wspubsub.UUID]*wspubsub.Client)
119 | 			}
120 | 
121 | 			if _, ok := clients[channel][client.ID()]; !ok {
122 | 				numClientsInChannels++
123 | 			}
124 | 
125 | 			clients[channel][client.ID()] = client
126 | 		}
127 | 	}
128 | 
129 | 	t.Run("Count clients in unknown channel", func(t *testing.T) {
130 | 		require.Equal(t, 0, clientStore.Count("UNKNOWN"))
131 | 	})
132 | 
133 | 	t.Run("Count clients in a channel", func(t *testing.T) {
134 | 		for _, channel := range availableChannels {
135 | 			require.Equal(t, len(clients[channel]), clientStore.Count(channel))
136 | 		}
137 | 	})
138 | 
139 | 	t.Run("Count clients in a few random channels", func(t *testing.T) {
140 | 		channel1 := availableChannels[rand.Intn(len(availableChannels))]
141 | 		channel2 := availableChannels[rand.Intn(len(availableChannels))]
142 | 		require.Equal(t, len(clients[channel1])+len(clients[channel2]), clientStore.Count(channel1, channel2))
143 | 	})
144 | 
145 | 	t.Run("Count clients in all channels", func(t *testing.T) {
146 | 		require.Equal(t, numClientsInChannels, clientStore.Count(availableChannels...))
147 | 	})
148 | 
149 | 	t.Run("Count clients without channels", func(t *testing.T) {
150 | 		require.Equal(t, numClients, clientStore.Count())
151 | 	})
152 | 
153 | 	t.Run("Find clients in unknown channel", func(t *testing.T) {
154 | 		foundClients := 0
155 | 		channel := "UNKNOWN"
156 | 
157 | 		fn := func(client wspubsub.WebsocketClient) error {
158 | 			foundClients++
159 | 
160 | 			return nil
161 | 		}
162 | 
163 | 		err := clientStore.Find(fn, channel)
164 | 		require.NoError(t, err)
165 | 
166 | 		require.Equal(t, 0, foundClients)
167 | 	})
168 | 
169 | 	t.Run("Find clients in a random channel", func(t *testing.T) {
170 | 		foundClients := 0
171 | 		channel := availableChannels[rand.Intn(len(availableChannels))]
172 | 
173 | 		fn := func(client wspubsub.WebsocketClient) error {
174 | 			foundClients++
175 | 			_, ok := clients[channel][client.ID()]
176 | 			require.True(t, ok)
177 | 
178 | 			return nil
179 | 		}
180 | 
181 | 		err := clientStore.Find(fn, channel)
182 | 		require.NoError(t, err)
183 | 
184 | 		require.Equal(t, len(clients[channel]), foundClients)
185 | 	})
186 | 
187 | 	t.Run("Find clients in a few random channels", func(t *testing.T) {
188 | 		numFoundClients := 0
189 | 		channel1 := availableChannels[rand.Intn(len(availableChannels))]
190 | 		channel2 := availableChannels[rand.Intn(len(availableChannels))]
191 | 
192 | 		fn := func(client wspubsub.WebsocketClient) error {
193 | 			numFoundClients++
194 | 			_, ok1 := clients[channel1][client.ID()]
195 | 			_, ok2 := clients[channel2][client.ID()]
196 | 
197 | 			require.True(t, ok1 || ok2)
198 | 
199 | 			return nil
200 | 		}
201 | 
202 | 		err := clientStore.Find(fn, channel1, channel2)
203 | 		require.NoError(t, err)
204 | 
205 | 		require.Equal(t, len(clients[channel1])+len(clients[channel2]), numFoundClients)
206 | 	})
207 | 
208 | 	t.Run("Find clients in all channels", func(t *testing.T) {
209 | 		numFoundClients := 0
210 | 		fn := func(client wspubsub.WebsocketClient) error {
211 | 			numFoundClients++
212 | 			exists := false
213 | 			for _, channel := range availableChannels {
214 | 				if _, ok := clients[channel][client.ID()]; ok {
215 | 					exists = true
216 | 					break
217 | 				}
218 | 			}
219 | 
220 | 			require.True(t, exists)
221 | 
222 | 			return nil
223 | 		}
224 | 
225 | 		err := clientStore.Find(fn, availableChannels...)
226 | 		require.NoError(t, err)
227 | 
228 | 		require.Equal(t, numClientsInChannels, numFoundClients)
229 | 	})
230 | 
231 | 	t.Run("Find clients without channels", func(t *testing.T) {
232 | 		foundClients := 0
233 | 		fn := func(client wspubsub.WebsocketClient) error {
234 | 			foundClients++
235 | 
236 | 			exists := false
237 | 			for _, channel := range availableChannels {
238 | 				if _, ok := clients[channel][client.ID()]; ok {
239 | 					exists = true
240 | 					break
241 | 				}
242 | 			}
243 | 
244 | 			require.True(t, exists)
245 | 
246 | 			return nil
247 | 		}
248 | 
249 | 		err := clientStore.Find(fn)
250 | 		require.NoError(t, err)
251 | 
252 | 		require.Equal(t, numClients, foundClients)
253 | 	})
254 | 
255 | 	t.Run("Find clients with iteration error", func(t *testing.T) {
256 | 		fn := func(client wspubsub.WebsocketClient) error {
257 | 			return errors.New("something went wrong")
258 | 		}
259 | 
260 | 		err := clientStore.Find(fn)
261 | 		require.Error(t, err)
262 | 	})
263 | }
264 | 
265 | func TestClientStore_Channels(t *testing.T) {
266 | 	ctrl := gomock.NewController(t)
267 | 	defer ctrl.Finish()
268 | 
269 | 	logger := mock.NewMockLogger(ctrl)
270 | 	upgrader := mock.NewMockWebsocketConnectionUpgrader(ctrl)
271 | 
272 | 	clientStoreOptions := wspubsub.NewClientStoreOptions()
273 | 	clientStore := wspubsub.NewClientStore(clientStoreOptions, logger)
274 | 
275 | 	clientOptions := wspubsub.NewClientOptions()
276 | 	client := wspubsub.NewClient(clientOptions, clientID, upgrader, logger)
277 | 	clientStore.Set(client)
278 | 
279 | 	unknownClientID := wspubsub.UUID{}
280 | 
281 | 	t.Run("Use channels for unknown client", func(t *testing.T) {
282 | 		channels, err := clientStore.Channels(unknownClientID)
283 | 		require.Empty(t, channels)
284 | 		require.Equal(t, wspubsub.NewClientNotFoundError(unknownClientID), errors.Cause(err).(*wspubsub.ClientNotFoundError))
285 | 
286 | 		numChannels, err := clientStore.CountChannels(unknownClientID)
287 | 		require.Equal(t, 0, numChannels)
288 | 		require.Equal(t, wspubsub.NewClientNotFoundError(unknownClientID), errors.Cause(err).(*wspubsub.ClientNotFoundError))
289 | 
290 | 		err = clientStore.SetChannels(unknownClientID, "X")
291 | 		require.Equal(t, wspubsub.NewClientNotFoundError(unknownClientID), errors.Cause(err).(*wspubsub.ClientNotFoundError))
292 | 
293 | 		err = clientStore.UnsetChannels(unknownClientID, "X")
294 | 		require.Equal(t, wspubsub.NewClientNotFoundError(unknownClientID), errors.Cause(err).(*wspubsub.ClientNotFoundError))
295 | 	})
296 | 
297 | 	t.Run("Check default sate of channels", func(t *testing.T) {
298 | 		channels, err := clientStore.Channels(clientID)
299 | 		require.NoError(t, err)
300 | 		require.Empty(t, channels)
301 | 		numChannels, err := clientStore.CountChannels(clientID)
302 | 		require.NoError(t, err)
303 | 		require.Equal(t, 0, numChannels)
304 | 	})
305 | 
306 | 	t.Run("Set empty channels", func(t *testing.T) {
307 | 		err := clientStore.SetChannels(clientID)
308 | 		require.NoError(t, err)
309 | 	})
310 | 
311 | 	t.Run("Set channels", func(t *testing.T) {
312 | 		channels := []string{"X", "Y", "W", "Z"}
313 | 		sort.Strings(channels)
314 | 
315 | 		err := clientStore.SetChannels(clientID, channels...)
316 | 		require.NoError(t, err)
317 | 		newChannels, err := clientStore.Channels(clientID)
318 | 		sort.Strings(newChannels)
319 | 		require.NoError(t, err)
320 | 		require.Equal(t, channels, newChannels)
321 | 		numChannels, err := clientStore.CountChannels(clientID)
322 | 		require.NoError(t, err)
323 | 		require.Equal(t, len(channels), numChannels)
324 | 	})
325 | 
326 | 	t.Run("Unset channels", func(t *testing.T) {
327 | 		channels := []string{"X", "Z"}
328 | 		sort.Strings(channels)
329 | 
330 | 		err := clientStore.UnsetChannels(clientID, "Y", "W")
331 | 		require.NoError(t, err)
332 | 		newChannels, err := clientStore.Channels(clientID)
333 | 		sort.Strings(newChannels)
334 | 		require.NoError(t, err)
335 | 		require.Equal(t, channels, newChannels)
336 | 		numChannels, err := clientStore.CountChannels(clientID)
337 | 		require.NoError(t, err)
338 | 		require.Equal(t, len(channels), numChannels)
339 | 
340 | 		err = clientStore.UnsetChannels(clientID)
341 | 		require.NoError(t, err)
342 | 		newChannels, err = clientStore.Channels(clientID)
343 | 		sort.Strings(newChannels)
344 | 		require.NoError(t, err)
345 | 		require.Empty(t, newChannels)
346 | 	})
347 | }
348 | 


--------------------------------------------------------------------------------
/client_test.go:
--------------------------------------------------------------------------------
  1 | package wspubsub_test
  2 | 
  3 | import (
  4 | 	"fmt"
  5 | 	"net/http/httptest"
  6 | 	"strings"
  7 | 	"testing"
  8 | 	"time"
  9 | 
 10 | 	"github.com/golang/mock/gomock"
 11 | 	"github.com/kpeu3i/wspubsub"
 12 | 	"github.com/kpeu3i/wspubsub/mock"
 13 | 	"github.com/pkg/errors"
 14 | 	"github.com/stretchr/testify/require"
 15 | )
 16 | 
 17 | var clientID = wspubsub.UUID([16]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16})
 18 | 
 19 | func TestClient_ID(t *testing.T) {
 20 | 	ctrl := gomock.NewController(t)
 21 | 	defer ctrl.Finish()
 22 | 
 23 | 	options := wspubsub.NewClientOptions()
 24 | 	upgrader := mock.NewMockWebsocketConnectionUpgrader(ctrl)
 25 | 	logger := mock.NewMockLogger(ctrl)
 26 | 	client := wspubsub.NewClient(options, clientID, upgrader, logger)
 27 | 
 28 | 	require.Equal(t, client.ID(), clientID)
 29 | }
 30 | 
 31 | func TestClient_ReuseConnectAndClose(t *testing.T) {
 32 | 	ctrl := gomock.NewController(t)
 33 | 	defer func() {
 34 | 		time.Sleep(100 * time.Millisecond)
 35 | 		ctrl.Finish()
 36 | 	}()
 37 | 
 38 | 	upgradeErr := errors.New("can't upgrade connection")
 39 | 
 40 | 	request1 := httptest.NewRequest("GET", "/", nil)
 41 | 	response1 := httptest.NewRecorder()
 42 | 
 43 | 	request2 := httptest.NewRequest("GET", "/", nil)
 44 | 	response2 := httptest.NewRecorder()
 45 | 
 46 | 	request3 := httptest.NewRequest("GET", "/", nil)
 47 | 	response3 := httptest.NewRecorder()
 48 | 
 49 | 	logger := mock.NewMockLogger(ctrl)
 50 | 
 51 | 	connection := mock.NewMockWebsocketConnection(ctrl)
 52 | 
 53 | 	connection.
 54 | 		EXPECT().
 55 | 		Read().
 56 | 		Times(2).
 57 | 		Do(func() {
 58 | 			time.Sleep(5 * time.Second)
 59 | 		})
 60 | 
 61 | 	connection.
 62 | 		EXPECT().
 63 | 		Close().
 64 | 		Times(2).
 65 | 		Return(nil)
 66 | 
 67 | 	upgrader := mock.NewMockWebsocketConnectionUpgrader(ctrl)
 68 | 
 69 | 	upgrader.
 70 | 		EXPECT().
 71 | 		Upgrade(gomock.Eq(response1), gomock.Eq(request1)).
 72 | 		Return(connection, nil).
 73 | 		Times(1)
 74 | 
 75 | 	upgrader.
 76 | 		EXPECT().
 77 | 		Upgrade(gomock.Eq(response2), gomock.Eq(request2)).
 78 | 		Return(connection, nil).
 79 | 		Times(1)
 80 | 
 81 | 	upgrader.
 82 | 		EXPECT().
 83 | 		Upgrade(gomock.Eq(response3), gomock.Eq(request3)).
 84 | 		Return(nil, upgradeErr).
 85 | 		Times(1)
 86 | 
 87 | 	options := wspubsub.NewClientOptions()
 88 | 	client := wspubsub.NewClient(options, clientID, upgrader, logger)
 89 | 
 90 | 	for i := 1; i <= 2; i++ {
 91 | 		t.Run(fmt.Sprintf("Connection attempt #%d", i), func(t *testing.T) {
 92 | 			err := client.Connect(response1, request1)
 93 | 			require.NoError(t, err)
 94 | 
 95 | 			err = client.Connect(response1, request1)
 96 | 			require.Error(t, err)
 97 | 			require.Equal(t, wspubsub.NewClientRepeatConnectError(clientID), errors.Cause(err).(*wspubsub.ClientRepeatConnectError))
 98 | 		})
 99 | 
100 | 		t.Run(fmt.Sprintf("Closing attempt #%d", i), func(t *testing.T) {
101 | 			err := client.Close()
102 | 			require.NoError(t, err)
103 | 
104 | 			err = client.Close()
105 | 			require.NoError(t, err)
106 | 		})
107 | 	}
108 | 
109 | 	t.Run("Connection upgrade error", func(t *testing.T) {
110 | 		err := client.Connect(response3, request3)
111 | 		require.Equal(t, wspubsub.NewClientConnectError(clientID, upgradeErr), errors.Cause(err).(*wspubsub.ClientConnectError))
112 | 	})
113 | }
114 | 
115 | func TestClient_CloseError(t *testing.T) {
116 | 	ctrl := gomock.NewController(t)
117 | 	defer func() {
118 | 		time.Sleep(100 * time.Millisecond)
119 | 		ctrl.Finish()
120 | 	}()
121 | 
122 | 	closeErrText := "close_error"
123 | 
124 | 	request := httptest.NewRequest("GET", "/", nil)
125 | 	response := httptest.NewRecorder()
126 | 
127 | 	logger := mock.NewMockLogger(ctrl)
128 | 
129 | 	connection := mock.NewMockWebsocketConnection(ctrl)
130 | 
131 | 	connection.
132 | 		EXPECT().
133 | 		Read().
134 | 		Times(1).
135 | 		Do(func() {
136 | 			time.Sleep(5 * time.Second)
137 | 		})
138 | 
139 | 	connection.
140 | 		EXPECT().
141 | 		Close().
142 | 		Times(1).
143 | 		Return(errors.New(closeErrText))
144 | 
145 | 	upgrader := mock.NewMockWebsocketConnectionUpgrader(ctrl)
146 | 
147 | 	upgrader.
148 | 		EXPECT().
149 | 		Upgrade(gomock.Eq(response), gomock.Eq(request)).
150 | 		Return(connection, nil).
151 | 		Times(1)
152 | 
153 | 	options := wspubsub.NewClientOptions()
154 | 	client := wspubsub.NewClient(options, clientID, upgrader, logger)
155 | 
156 | 	err := client.Connect(response, request)
157 | 	require.NoError(t, err)
158 | 
159 | 	err = client.Close()
160 | 	require.Error(t, err)
161 | 	require.True(t, strings.Contains(err.Error(), closeErrText))
162 | }
163 | 
164 | func TestClient_Write(t *testing.T) {
165 | 	ctrl := gomock.NewController(t)
166 | 	defer func() {
167 | 		time.Sleep(100 * time.Millisecond)
168 | 		ctrl.Finish()
169 | 	}()
170 | 
171 | 	message := wspubsub.NewTextMessageFromString("TEST")
172 | 	closedErr := wspubsub.NewConnectionClosedError(errors.New("i/o timeout"))
173 | 
174 | 	request := httptest.NewRequest("GET", "/", nil)
175 | 	response := httptest.NewRecorder()
176 | 
177 | 	logger := mock.NewMockLogger(ctrl)
178 | 
179 | 	connection := mock.NewMockWebsocketConnection(ctrl)
180 | 	connection.
181 | 		EXPECT().
182 | 		Read().
183 | 		Times(1).
184 | 		Do(func() {
185 | 			time.Sleep(5 * time.Second)
186 | 		})
187 | 
188 | 	connection.
189 | 		EXPECT().
190 | 		Write(gomock.Eq(message)).
191 | 		Times(1)
192 | 
193 | 	connection.
194 | 		EXPECT().
195 | 		Write(gomock.Eq(message)).
196 | 		Times(1).
197 | 		Return(closedErr)
198 | 
199 | 	upgrader := mock.NewMockWebsocketConnectionUpgrader(ctrl)
200 | 	upgrader.
201 | 		EXPECT().
202 | 		Upgrade(gomock.Eq(response), gomock.Eq(request)).
203 | 		Return(connection, nil).
204 | 		Times(1)
205 | 
206 | 	options := wspubsub.NewClientOptions()
207 | 	client := wspubsub.NewClient(options, clientID, upgrader, logger)
208 | 	client.OnError(func(id wspubsub.UUID, err error) {
209 | 		require.Equal(t, clientID, id)
210 | 		require.Equal(t, wspubsub.NewClientSendError(clientID, message, closedErr), errors.Cause(err).(*wspubsub.ClientSendError))
211 | 	})
212 | 
213 | 	err := client.Connect(response, request)
214 | 	require.NoError(t, err)
215 | 
216 | 	err = client.Send(message)
217 | 	require.NoError(t, err)
218 | 
219 | 	err = client.Send(message)
220 | 	require.NoError(t, err)
221 | }
222 | 
223 | func TestClient_WriteBufferOverflow(t *testing.T) {
224 | 	ctrl := gomock.NewController(t)
225 | 	defer func() {
226 | 		time.Sleep(100 * time.Millisecond)
227 | 		ctrl.Finish()
228 | 	}()
229 | 
230 | 	closedErr := wspubsub.NewConnectionClosedError(errors.New("i/o timeout"))
231 | 	message := wspubsub.NewTextMessageFromString("TEST")
232 | 
233 | 	logger := mock.NewMockLogger(ctrl)
234 | 	upgrader := mock.NewMockWebsocketConnectionUpgrader(ctrl)
235 | 
236 | 	options := wspubsub.NewClientOptions()
237 | 	options.SendBufferSize = 1
238 | 	client := wspubsub.NewClient(options, clientID, upgrader, logger)
239 | 	client.OnError(func(id wspubsub.UUID, err error) {
240 | 		require.Equal(t, clientID, id)
241 | 		require.Equal(t, wspubsub.NewClientSendError(clientID, message, closedErr), errors.Cause(err).(*wspubsub.ClientSendError))
242 | 	})
243 | 
244 | 	err := client.Send(message)
245 | 	require.NoError(t, err)
246 | 
247 | 	err = client.Send(message)
248 | 	require.Equal(t, wspubsub.NewClientSendBufferOverflowError(clientID), errors.Cause(err).(*wspubsub.ClientSendBufferOverflowError))
249 | }
250 | 
251 | func TestClient_Ping(t *testing.T) {
252 | 	ctrl := gomock.NewController(t)
253 | 	defer func() {
254 | 		time.Sleep(100 * time.Millisecond)
255 | 		ctrl.Finish()
256 | 	}()
257 | 
258 | 	message := wspubsub.NewPingMessage()
259 | 	closedErr := wspubsub.NewConnectionClosedError(errors.New("i/o timeout"))
260 | 
261 | 	request := httptest.NewRequest("GET", "/", nil)
262 | 	response := httptest.NewRecorder()
263 | 
264 | 	options := wspubsub.NewClientOptions()
265 | 	options.PingInterval = 1 * time.Millisecond
266 | 
267 | 	logger := mock.NewMockLogger(ctrl)
268 | 
269 | 	connection := mock.NewMockWebsocketConnection(ctrl)
270 | 
271 | 	connection.
272 | 		EXPECT().
273 | 		Read().
274 | 		Times(1).
275 | 		Do(func() {
276 | 			time.Sleep(5 * time.Second)
277 | 		})
278 | 
279 | 	connection.
280 | 		EXPECT().
281 | 		Write(gomock.Eq(message)).
282 | 		Times(1).
283 | 		Return(nil)
284 | 
285 | 	connection.
286 | 		EXPECT().
287 | 		Write(gomock.Eq(message)).
288 | 		Times(1).
289 | 		Return(closedErr)
290 | 
291 | 	upgrader := mock.NewMockWebsocketConnectionUpgrader(ctrl)
292 | 
293 | 	upgrader.
294 | 		EXPECT().
295 | 		Upgrade(gomock.Eq(response), gomock.Eq(request)).
296 | 		Return(connection, nil).
297 | 		Times(1)
298 | 
299 | 	client := wspubsub.NewClient(options, clientID, upgrader, logger)
300 | 	client.OnError(func(id wspubsub.UUID, err error) {
301 | 		require.Equal(t, clientID, id)
302 | 		require.Equal(t, wspubsub.NewClientPingError(clientID, message, closedErr), errors.Cause(err).(*wspubsub.ClientPingError))
303 | 	})
304 | 
305 | 	err := client.Connect(response, request)
306 | 	require.NoError(t, err)
307 | }
308 | 
309 | func TestClient_Read(t *testing.T) {
310 | 	ctrl := gomock.NewController(t)
311 | 	defer func() {
312 | 		time.Sleep(100 * time.Millisecond)
313 | 		ctrl.Finish()
314 | 	}()
315 | 
316 | 	message := wspubsub.NewTextMessageFromString("TEST")
317 | 	closedErr := wspubsub.NewConnectionClosedError(errors.New("i/o timeout"))
318 | 
319 | 	request := httptest.NewRequest("GET", "/", nil)
320 | 	response := httptest.NewRecorder()
321 | 
322 | 	logger := mock.NewMockLogger(ctrl)
323 | 
324 | 	connection := mock.NewMockWebsocketConnection(ctrl)
325 | 
326 | 	connection.
327 | 		EXPECT().
328 | 		Read().
329 | 		Times(1).
330 | 		Return(message, nil)
331 | 
332 | 	connection.
333 | 		EXPECT().
334 | 		Read().
335 | 		Times(1).
336 | 		Return(message, closedErr)
337 | 
338 | 	upgrader := mock.NewMockWebsocketConnectionUpgrader(ctrl)
339 | 
340 | 	upgrader.
341 | 		EXPECT().
342 | 		Upgrade(gomock.Eq(response), gomock.Eq(request)).
343 | 		Return(connection, nil).
344 | 		Times(1)
345 | 
346 | 	options := wspubsub.NewClientOptions()
347 | 	client := wspubsub.NewClient(options, clientID, upgrader, logger)
348 | 
349 | 	client.OnReceive(func(id wspubsub.UUID, message wspubsub.Message) {
350 | 		require.Equal(t, clientID, id)
351 | 		require.Equal(t, message, message)
352 | 	})
353 | 
354 | 	client.OnError(func(id wspubsub.UUID, err error) {
355 | 		require.Equal(t, clientID, id)
356 | 		require.Equal(t, wspubsub.NewClientReceiveError(clientID, message, closedErr), errors.Cause(err).(*wspubsub.ClientReceiveError))
357 | 	})
358 | 
359 | 	err := client.Connect(response, request)
360 | 	require.NoError(t, err)
361 | }
362 | 


--------------------------------------------------------------------------------
/connection_closed_error.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import (
 4 | 	"fmt"
 5 | 
 6 | 	"github.com/pkg/errors"
 7 | )
 8 | 
 9 | // ConnectionClosedError returned when trying to read or write to closed WebSocket connection.
10 | type ConnectionClosedError struct {
11 | 	Err error
12 | }
13 | 
14 | // ConnectionClosedError implements an error interface.
15 | func (e *ConnectionClosedError) Error() string {
16 | 	return fmt.Sprintf("wspubsub: connection is closed: err=%s", e.Err)
17 | }
18 | 
19 | // NewConnectionClosedError initializes a new ConnectionClosedError.
20 | func NewConnectionClosedError(err error) *ConnectionClosedError {
21 | 	return &ConnectionClosedError{Err: err}
22 | }
23 | 
24 | // IsConnectionClosedError checks if error type is ConnectionClosedError.
25 | func IsConnectionClosedError(err error) (*ConnectionClosedError, bool) {
26 | 	v, ok := errors.Cause(err).(*ConnectionClosedError)
27 | 
28 | 	return v, ok
29 | }
30 | 


--------------------------------------------------------------------------------
/connection_closed_error_test.go:
--------------------------------------------------------------------------------
 1 | package wspubsub_test
 2 | 
 3 | import (
 4 | 	"errors"
 5 | 	"testing"
 6 | 
 7 | 	"github.com/kpeu3i/wspubsub"
 8 | 	"github.com/stretchr/testify/require"
 9 | )
10 | 
11 | func TestConnectionClosedError(t *testing.T) {
12 | 	rawErr := errors.New("TEST")
13 | 	err := wspubsub.NewConnectionClosedError(rawErr)
14 | 	require.Equal(t, rawErr, err.Err)
15 | 	require.NotEmpty(t, clientID, err.Error())
16 | 
17 | 	e, ok := wspubsub.IsConnectionClosedError(err)
18 | 	require.NotNil(t, e)
19 | 	require.True(t, ok)
20 | 
21 | 	e, ok = wspubsub.IsConnectionClosedError(rawErr)
22 | 	require.Nil(t, e)
23 | 	require.False(t, ok)
24 | }
25 | 


--------------------------------------------------------------------------------
/doc.go:
--------------------------------------------------------------------------------
1 | // Package wspubsub provides an easy way to publish/subscribe and receive messages
2 | // over WebSocket protocol hiding the details of the underlying transport.
3 | // The library based on the idea that all messages are published to channels.
4 | // Subscribers will receive messages published to the channels to which they subscribe.
5 | // The hub is responsible for defining the channels to which subscribers can subscribe.
6 | package wspubsub
7 | 


--------------------------------------------------------------------------------
/examples/bench/client/main.go:
--------------------------------------------------------------------------------
  1 | package main
  2 | 
  3 | import (
  4 | 	"flag"
  5 | 	"fmt"
  6 | 	"math"
  7 | 	"math/rand"
  8 | 	"strings"
  9 | 	"sync"
 10 | 	"time"
 11 | 
 12 | 	"github.com/gorilla/websocket"
 13 | 	"github.com/kpeu3i/wspubsub"
 14 | )
 15 | 
 16 | var (
 17 | 	logger wspubsub.Logger
 18 | )
 19 | 
 20 | func init() {
 21 | 	rand.Seed(time.Now().Unix())
 22 | 	logger = wspubsub.NewLogrusLogger(wspubsub.NewLogrusLoggerOptions())
 23 | }
 24 | 
 25 | func main() {
 26 | 	url := flag.String("addr", "ws://localhost:8080/", "Server URL")
 27 | 	connectionCount := flag.Int("connections", 5000, "Total connection count")
 28 | 	workerCount := flag.Int("workers", 1, "Total worker count")
 29 | 	channels := flag.String("channels", "X,Y,Z", "One of possible subscription channels will be used")
 30 | 	wait := flag.String("wait", "10s", "Wait time after all connections will be established")
 31 | 	readTimeout := flag.String("read_timeout", "60s", "Read timeout")
 32 | 	writeTimeout := flag.String("write_timeout", "10s", "Write timeout")
 33 | 
 34 | 	channelList := strings.Split(*channels, ",")
 35 | 
 36 | 	readTimeoutDuration, err := time.ParseDuration(*readTimeout)
 37 | 	if err != nil {
 38 | 		panic(err)
 39 | 	}
 40 | 
 41 | 	writeTimeoutDuration, err := time.ParseDuration(*writeTimeout)
 42 | 	if err != nil {
 43 | 		panic(err)
 44 | 	}
 45 | 
 46 | 	waitDuration, err := time.ParseDuration(*wait)
 47 | 	if err != nil {
 48 | 		panic(err)
 49 | 	}
 50 | 
 51 | 	wg := sync.WaitGroup{}
 52 | 	connectionsMutex := sync.RWMutex{}
 53 | 	connections := make([]*websocket.Conn, 0, *connectionCount)
 54 | 	tasks := make(chan func() error)
 55 | 
 56 | 	logger.Println("Starting...")
 57 | 
 58 | 	for i := 1; i <= *workerCount; i++ {
 59 | 		go runWorker(tasks)
 60 | 	}
 61 | 
 62 | 	now := time.Now()
 63 | 
 64 | 	for i := 1; i <= *connectionCount; i++ {
 65 | 		wg.Add(1)
 66 | 		tasks <- func() error {
 67 | 			defer wg.Done()
 68 | 
 69 | 			connection, err := subscribe(*url, channelList, readTimeoutDuration, writeTimeoutDuration)
 70 | 			if err != nil {
 71 | 				return err
 72 | 			}
 73 | 
 74 | 			connectionsMutex.Lock()
 75 | 			connections = append(connections, connection)
 76 | 			connectionsMutex.Unlock()
 77 | 
 78 | 			if i%100 == 0 {
 79 | 				logger.Printf("Connected: %d\n", i)
 80 | 			}
 81 | 
 82 | 			return nil
 83 | 		}
 84 | 	}
 85 | 
 86 | 	wg.Wait()
 87 | 
 88 | 	took := time.Since(now)
 89 | 
 90 | 	logger.Println("Done!")
 91 | 	logger.Printf("Waiting %s...", waitDuration)
 92 | 
 93 | 	time.Sleep(waitDuration)
 94 | 
 95 | 	logger.Println("Stopping...")
 96 | 
 97 | 	close(tasks)
 98 | 
 99 | 	for _, connection := range connections {
100 | 		_ = connection.Close()
101 | 	}
102 | 
103 | 	logger.Println("Connections:", len(connections))
104 | 	logger.Println("Took:", took)
105 | 	logger.Println("RPS:", math.Round(float64(len(connections))/took.Seconds()))
106 | }
107 | 
108 | func subscribe(url string, channels []string, readTimeout, writeTimeout time.Duration) (*websocket.Conn, error) {
109 | 	connection, _, err := websocket.DefaultDialer.Dial(url, nil)
110 | 	if err != nil {
111 | 		return nil, err
112 | 	}
113 | 
114 | 	err = connection.SetReadDeadline(time.Now().Add(readTimeout))
115 | 	if err != nil {
116 | 		return nil, err
117 | 	}
118 | 
119 | 	err = connection.SetWriteDeadline(time.Now().Add(writeTimeout))
120 | 	if err != nil {
121 | 		return nil, err
122 | 	}
123 | 
124 | 	connection.SetPongHandler(func(string) error {
125 | 		return connection.SetReadDeadline(time.Now().Add(readTimeout))
126 | 	})
127 | 
128 | 	err = connection.WriteMessage(websocket.TextMessage, subscribeMessagePayload(channels))
129 | 	if err != nil {
130 | 		return nil, err
131 | 	}
132 | 
133 | 	// Reader
134 | 	go func() {
135 | 		for {
136 | 			err = connection.SetReadDeadline(time.Now().Add(60 * time.Second))
137 | 			if err != nil {
138 | 				if !strings.Contains(err.Error(), "use of closed network connection") {
139 | 					logger.Errorf("Read error: %s\n", err)
140 | 					return
141 | 				}
142 | 
143 | 				return
144 | 			}
145 | 
146 | 			_, _, err := connection.ReadMessage()
147 | 			if err != nil {
148 | 				if !strings.Contains(err.Error(), "use of closed network connection") {
149 | 					logger.Errorf("Read error: %s\n", err)
150 | 					return
151 | 				}
152 | 
153 | 				return
154 | 			}
155 | 		}
156 | 	}()
157 | 
158 | 	return connection, nil
159 | }
160 | 
161 | func runWorker(tasks <-chan func() error) {
162 | 	for task := range tasks {
163 | 		err := task()
164 | 		if err != nil {
165 | 			logger.Errorf("Failed to process a task: %s\n", err)
166 | 		}
167 | 	}
168 | }
169 | 
170 | func subscribeMessagePayload(channels []string) []byte {
171 | 	channel := channels[rand.Intn(len(channels))]
172 | 	message := fmt.Sprintf(`{"command": "SUBSCRIBE", "channels": ["%s"]}`, channel)
173 | 
174 | 	return []byte(message)
175 | }
176 | 


--------------------------------------------------------------------------------
/examples/bench/server/main.go:
--------------------------------------------------------------------------------
  1 | package main
  2 | 
  3 | import (
  4 | 	"encoding/json"
  5 | 	"expvar"
  6 | 	"flag"
  7 | 	"math"
  8 | 	"math/rand"
  9 | 	"net/http"
 10 | 	"os"
 11 | 	"os/signal"
 12 | 	"runtime"
 13 | 	"strings"
 14 | 	"syscall"
 15 | 	"time"
 16 | 
 17 | 	"github.com/kpeu3i/wspubsub"
 18 | 	"github.com/pkg/errors"
 19 | )
 20 | 
 21 | var (
 22 | 	publishMessageCount int64
 23 | 	publishTime         int64
 24 | 	publishCount        int64
 25 | )
 26 | 
 27 | type Message struct {
 28 | 	Command  string   `json:"command"`
 29 | 	Channels []string `json:"channels"`
 30 | }
 31 | 
 32 | func init() {
 33 | 	rand.Seed(time.Now().Unix())
 34 | 	expvar.Publish("Goroutines", expvar.Func(func() interface{} {
 35 | 		return runtime.NumGoroutine()
 36 | 	}))
 37 | }
 38 | 
 39 | func main() {
 40 | 	addr := flag.String("addr", "localhost:8080", "Server host and port")
 41 | 	path := flag.String("path", "/", "URL")
 42 | 	publishInterval := flag.String("publish", "50ms", "Messages publishing interval")
 43 | 	publishChannels := flag.String("channels", "X,Y,Z", "Publishing channels")
 44 | 	messageSize := flag.Int("message", 5*1024, "Message size (bytes)")
 45 | 
 46 | 	publishChannelList := strings.Split(*publishChannels, ",")
 47 | 	publishIntervalDuration, err := time.ParseDuration(*publishInterval)
 48 | 	if err != nil {
 49 | 		panic(err)
 50 | 	}
 51 | 
 52 | 	publishTicker := time.NewTicker(publishIntervalDuration)
 53 | 	payload := messagePayload(*messageSize)
 54 | 
 55 | 	quit := make(chan os.Signal, 1)
 56 | 	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
 57 | 
 58 | 	hub := wspubsub.NewDefaultHub()
 59 | 
 60 | 	hub.OnReceive(func(clientID wspubsub.UUID, message wspubsub.Message) {
 61 | 		m := Message{}
 62 | 		err := json.Unmarshal(message.Payload, &m)
 63 | 		if err != nil {
 64 | 			hub.LogError(errors.Wrap(err, "hub.on_receive.unmarshal"))
 65 | 			return
 66 | 		}
 67 | 
 68 | 		switch m.Command {
 69 | 		case "SUBSCRIBE":
 70 | 			err := hub.Subscribe(clientID, m.Channels...)
 71 | 			if err != nil {
 72 | 				hub.LogError(errors.Wrap(err, "hub.on_receive.subscribe"))
 73 | 			}
 74 | 		case "UNSUBSCRIBE":
 75 | 			err := hub.Unsubscribe(clientID, m.Channels...)
 76 | 			if err != nil {
 77 | 				hub.LogError(errors.Wrap(err, "hub.on_receive.unsubscribe"))
 78 | 			}
 79 | 		}
 80 | 	})
 81 | 
 82 | 	http.Handle(*path, hub)
 83 | 
 84 | 	go func() {
 85 | 		err := http.ListenAndServe(*addr, nil)
 86 | 		if err != nil {
 87 | 			hub.LogPanic(err)
 88 | 		}
 89 | 	}()
 90 | 
 91 | 	go publish(hub, publishTicker, publishChannelList, payload)
 92 | 
 93 | 	<-quit
 94 | 
 95 | 	publishTicker.Stop()
 96 | 
 97 | 	err = hub.Close()
 98 | 	if err != nil {
 99 | 		hub.LogPanic(err)
100 | 	}
101 | 
102 | 	if publishTime > 0 && publishCount > 0 {
103 | 		hub.LogPrintf("Total publish time: %s\n", time.Duration(publishTime))
104 | 		hub.LogPrintf("Total messages published: %d\n", publishMessageCount)
105 | 		hub.LogPrintf("Single message publish time: %s\n", time.Duration(publishTime/publishCount))
106 | 		hub.LogPrintf("Publish RPS: %d\n", int(float64(publishMessageCount)/(float64(publishTime)/1e9)))
107 | 	}
108 | }
109 | 
110 | func publish(hub *wspubsub.Hub, ticker *time.Ticker, channels []string, payload []byte) {
111 | 	hub.LogInfoln("Starting publishing...")
112 | 	message := wspubsub.NewBinaryMessage(payload)
113 | 	for range ticker.C {
114 | 		channel := channels[rand.Intn(len(channels))]
115 | 		now := time.Now()
116 | 		clientCount, err := hub.Publish(message, channel)
117 | 		if err != nil {
118 | 			hub.LogPanic(err)
119 | 		}
120 | 
121 | 		if clientCount > 0 {
122 | 			publishMessageCount += int64(clientCount)
123 | 			publishTime += time.Since(now).Nanoseconds()
124 | 			publishCount++
125 | 		}
126 | 	}
127 | }
128 | 
129 | func messagePayload(size int) []byte {
130 | 	message := make([]byte, 0, size)
131 | 	for i := 0; i < size; i++ {
132 | 		message = append(message, byte(rand.Intn(math.MaxUint8+1)))
133 | 	}
134 | 
135 | 	return message
136 | }
137 | 


--------------------------------------------------------------------------------
/examples/demo/README.md:
--------------------------------------------------------------------------------
 1 | # Demo
 2 | 
 3 | This example shows how to use the library to implement a simple pub/sub application.
 4 | 
 5 | ## Running the example
 6 | 
 7 | The example requires a working Go development environment.
 8 | 
 9 | ### Server
10 | 
11 |     $ go run main.go
12 | 
13 | Next command line options are allowed:
14 | 
15 | Option | Default | Description
16 | --- | --- | ---
17 | addr | localhost:8080 | Server host and port
18 | path | / | Server URL
19 | publish | 1s | Messages publishing interval
20 | channels | general,public,private | Publishing channels
21 | 
22 | ### WEB
23 | 
24 | 1. Open `index.html` in your browser
25 | 2. Click **Connect** button
26 | 2. Click **Subscribe** button
27 | 
28 | Now, every published message which is matched your subscription will be shown in the **Messages** field.


--------------------------------------------------------------------------------
/examples/demo/index.html:
--------------------------------------------------------------------------------
  1 | <!doctype html>
  2 | <html lang="en">
  3 | <head>
  4 |     <!-- Required meta tags -->
  5 |     <meta charset="utf-8">
  6 |     <meta name="viewport" content="width=device-width, initial-scale=1, shrink-to-fit=no">
  7 |     <!-- Bootstrap CSS -->
  8 |     <link rel="stylesheet" href="https://maxcdn.bootstrapcdn.com/bootstrap/4.0.0/css/bootstrap.min.css" integrity="sha384-Gn5384xqQ1aoWXA+058RXPxPg6fy4IWvTNh0E263XmFcJlSAwiGgFAW/dAiS6JXm" crossorigin="anonymous">
  9 |     <title>WSPubSub</title>
 10 |     <script type="application/javascript">
 11 |         let connection;
 12 | 
 13 |         document.addEventListener("DOMContentLoaded", function () {
 14 |             document.getElementById("connect").addEventListener("click", function () {
 15 |                 let onOpen = function(event) {
 16 |                     document.getElementById("status").className = "badge badge-success";
 17 |                     document.getElementById("status").textContent = "Connected"
 18 |                 };
 19 | 
 20 |                 let onMessage = function(event) {
 21 |                     document.getElementById("messages").value += "--> " + event.data + "\n";
 22 |                     document.getElementById("messages").scrollTop = document.getElementById("messages").scrollHeight;
 23 |                 };
 24 | 
 25 |                 let onClose = function(event) {
 26 |                     document.getElementById("status").className = "badge badge-secondary";
 27 |                     document.getElementById("status").textContent = "Not connected"
 28 |                 };
 29 | 
 30 |                 let onError = function(error) {
 31 |                     document.getElementById("status").className = "badge badge-error";
 32 |                     document.getElementById("status").value = `Not connected (error: ${error.message})`
 33 |                 };
 34 | 
 35 |                 connection = connect(
 36 |                     document.getElementById("url").value,
 37 |                     onOpen,
 38 |                     onMessage,
 39 |                     onClose,
 40 |                     onError
 41 |                 )
 42 |             });
 43 | 
 44 |             document.getElementById("disconnect").addEventListener("click", function () {
 45 |                 disconnect(connection)
 46 |             });
 47 | 
 48 |             document.getElementById("subscribe").addEventListener("click", function () {
 49 |                 let channel = document.getElementById("channel").value;
 50 |                 let message = subscribe(connection, channel);
 51 | 
 52 |                 if (message !== "") {
 53 |                     document.getElementById("messages").value += "<-- " + message + "\n";
 54 | 
 55 |                     let found = false;
 56 |                     let children = document.getElementById("subscription").children;
 57 |                     for (let i = 0; i < children.length; i++) {
 58 |                         if (children[i].textContent === channel) {
 59 |                             found = true;
 60 |                             break;
 61 |                         }
 62 |                     }
 63 | 
 64 |                     if (!found) {
 65 |                         if (document.getElementById("subscription").textContent === "None") {
 66 |                             document.getElementById("subscription").textContent = ""
 67 |                         }
 68 | 
 69 |                         let link = document.createElement('a');
 70 |                         link.href = "#";
 71 |                         link.className = "badge badge-warning subscription-channel";
 72 |                         link.text = channel;
 73 |                         link.onclick = function() {
 74 |                             document.getElementById("channel").value = channel;
 75 | 
 76 |                             return false
 77 |                         };
 78 |                         document.getElementById("subscription").appendChild(link);
 79 |                         document.getElementById("subscription").appendChild(document.createTextNode(" "));
 80 |                     }
 81 |                 }
 82 |             });
 83 | 
 84 |             document.getElementById("unsubscribe").addEventListener("click", function () {
 85 |                 let channel = document.getElementById("channel").value;
 86 |                 let message = unsubscribe(connection, channel);
 87 | 
 88 |                 if (message !== "") {
 89 |                     document.getElementById("messages").value += "<-- " + message + "\n";
 90 | 
 91 |                     let children = document.getElementById("subscription").children;
 92 |                     for (let i = 0; i < children.length; i++) {
 93 |                         if (children[i].textContent === channel) {
 94 |                             children[i].remove();
 95 |                             break;
 96 |                         }
 97 |                     }
 98 | 
 99 |                     if (document.getElementById("subscription").textContent.trim() === "") {
100 |                         document.getElementById("subscription").textContent = "None"
101 |                     }
102 |                 }
103 |             });
104 |         });
105 | 
106 |         function connect(url, onOpen, onMessage, onClose, onError) {
107 |             let connection = new WebSocket(url);
108 | 
109 |             connection.onopen = function(event) {
110 |                 onOpen(event)
111 |             };
112 | 
113 |             connection.onmessage = function(event) {
114 |                 onMessage(event)
115 |             };
116 | 
117 |             connection.onclose = function(event) {
118 |                 onClose(event)
119 |             };
120 | 
121 |             connection.onerror = function(error) {
122 |                 onError(error)
123 |             };
124 | 
125 |             return connection
126 |         }
127 | 
128 |         function disconnect(connection) {
129 |             if (connection instanceof WebSocket) {
130 |                 if (connection.readyState === WebSocket.OPEN) {
131 |                     connection.close();
132 |                 }
133 |             }
134 |         }
135 | 
136 |         function subscribe(connection, channel) {
137 |             if (connection instanceof WebSocket) {
138 |                 if (connection.readyState === WebSocket.OPEN) {
139 |                     let message = `{"command": "SUBSCRIBE", "channels": ["${channel}"]}`;
140 |                     connection.send(message);
141 | 
142 |                     return message
143 |                 }
144 |             }
145 | 
146 |             return ""
147 |         }
148 | 
149 |         function unsubscribe(connection, channel) {
150 |             if (connection instanceof WebSocket) {
151 |                 if (connection.readyState === WebSocket.OPEN) {
152 |                     let message = `{"command": "UNSUBSCRIBE", "channels": ["${channel}"]}`;
153 |                     connection.send(message);
154 | 
155 |                     return message
156 |                 }
157 |             }
158 | 
159 |             return ""
160 |         }
161 |     </script>
162 | </head>
163 | <body>
164 |     <div class="container-fluid">
165 |         <main class="col-12 col-md-9 col-xl-8 py-md-3 pl-md-5" role="main">
166 |             <form>
167 |                 <div class="form-group">
168 |                     <div class="form-group">
169 |                         <label for="url">URL</label>
170 |                         <input type="text" class="form-control" id="url" placeholder="URL" value="ws://localhost:8080/">
171 |                     </div>
172 |                     <div class="form-group">
173 |                         Connection status:  <span id="status" class="badge badge-secondary">Not connected</span>
174 |                     </div>
175 |                     <button type="button" class="btn btn-primary" id="connect">Connect</button>
176 |                     <button type="button" class="btn btn-primary" id="disconnect">Disconnect</button>
177 | 
178 |                 </div>
179 |                 <div class="form-group">
180 |                     <div class="form-group">
181 |                         <label for="channel">Channel</label>
182 |                         <input type="text" class="form-control" id="channel" placeholder="Channel" value="general">
183 |                     </div>
184 |                     <div class="form-group">
185 |                         Subscription: <div class="d-inline" id="subscription">None</div>
186 |                     </div>
187 |                     <button type="button" class="btn btn-primary" id="subscribe">Subscribe</button>
188 |                     <button type="button" class="btn btn-primary" id="unsubscribe">Unsubscribe</button>
189 |                 </div>
190 |                 <div class="form-group">
191 |                     <label for="messages">Messages</label>
192 |                     <textarea class="form-control" id="messages" rows="10"></textarea>
193 |                 </div>
194 |             </form>
195 |         </main>
196 |     </div>
197 |     <!-- Optional JavaScript -->
198 |     <!-- jQuery first, then Popper.js, then Bootstrap JS -->
199 |     <script src="https://code.jquery.com/jquery-3.2.1.slim.min.js" integrity="sha384-KJ3o2DKtIkvYIK3UENzmM7KCkRr/rE9/Qpg6aAZGJwFDMVNA/GpGFF93hXpG5KkN" crossorigin="anonymous"></script>
200 |     <script src="https://cdnjs.cloudflare.com/ajax/libs/popper.js/1.12.9/umd/popper.min.js" integrity="sha384-ApNbgh9B+Y1QKtv3Rn7W3mgPxhU9K/ScQsAP7hUibX39j7fakFPskvXusvfa0b4Q" crossorigin="anonymous"></script>
201 |     <script src="https://maxcdn.bootstrapcdn.com/bootstrap/4.0.0/js/bootstrap.min.js" integrity="sha384-JZR6Spejh4U02d8jOt6vLEHfe/JQGiRRSQQxSfFWpi1MquVdAyjUar5+76PVCmYl" crossorigin="anonymous"></script>
202 | </body>
203 | </html>


--------------------------------------------------------------------------------
/examples/demo/main.go:
--------------------------------------------------------------------------------
 1 | package main
 2 | 
 3 | import (
 4 | 	"encoding/json"
 5 | 	"flag"
 6 | 	"fmt"
 7 | 	"math/rand"
 8 | 	"os"
 9 | 	"os/signal"
10 | 	"strings"
11 | 	"syscall"
12 | 	"time"
13 | 
14 | 	"github.com/kpeu3i/wspubsub"
15 | 	"github.com/pkg/errors"
16 | )
17 | 
18 | type Message struct {
19 | 	Command  string   `json:"command"`
20 | 	Channels []string `json:"channels"`
21 | }
22 | 
23 | func main() {
24 | 	addr := flag.String("addr", "localhost:8080", "Server host and port")
25 | 	path := flag.String("path", "/", "URL")
26 | 	publishInterval := flag.String("publish", "1s", "Messages publishing interval")
27 | 	publishChannels := flag.String("channels", "general,public,private", "Publishing channels")
28 | 
29 | 	publishChannelList := strings.Split(*publishChannels, ",")
30 | 	publishIntervalDuration, err := time.ParseDuration(*publishInterval)
31 | 	if err != nil {
32 | 		panic(err)
33 | 	}
34 | 
35 | 	rand.Seed(time.Now().Unix())
36 | 	quit := make(chan os.Signal, 1)
37 | 	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
38 | 
39 | 	hub := wspubsub.NewDefaultHub()
40 | 	defer hub.Close()
41 | 
42 | 	publishTicker := time.NewTicker(publishIntervalDuration)
43 | 	defer publishTicker.Stop()
44 | 
45 | 	hub.OnReceive(func(clientID wspubsub.UUID, message wspubsub.Message) {
46 | 		m := Message{}
47 | 		err := json.Unmarshal(message.Payload, &m)
48 | 		if err != nil {
49 | 			hub.LogError(errors.Wrap(err, "hub.on_receive.unmarshal"))
50 | 			return
51 | 		}
52 | 
53 | 		switch m.Command {
54 | 		case "SUBSCRIBE":
55 | 			err := hub.Subscribe(clientID, m.Channels...)
56 | 			if err != nil {
57 | 				hub.LogError(errors.Wrap(err, "hub.on_receive.subscribe"))
58 | 			}
59 | 			hub.LogInfof("Subscribed: client_id=%s, channels=%s", clientID, strings.Join(m.Channels, ","))
60 | 		case "UNSUBSCRIBE":
61 | 			err := hub.Unsubscribe(clientID, m.Channels...)
62 | 			if err != nil {
63 | 				hub.LogError(errors.Wrap(err, "hub.on_receive.unsubscribe"))
64 | 			}
65 | 			hub.LogInfof("Unsubscribed: client_id=%s, channels=%s", clientID, strings.Join(m.Channels, ","))
66 | 		}
67 | 	})
68 | 
69 | 	go func() {
70 | 		err := hub.ListenAndServe(*addr, *path)
71 | 		if err != nil {
72 | 			hub.LogPanic(err)
73 | 		}
74 | 	}()
75 | 
76 | 	go func() {
77 | 		for range publishTicker.C {
78 | 			// Pick a random channel
79 | 			channel := publishChannelList[rand.Intn(len(publishChannelList))]
80 | 			message := wspubsub.NewTextMessageFromString(fmt.Sprintf(`{"now": %d}`, time.Now().Unix()))
81 | 			_, err := hub.Publish(message, channel)
82 | 			if err != nil {
83 | 				hub.LogPanic(err)
84 | 			}
85 | 			hub.LogInfof("Published: channel=%s, message=%s", channel, string(message.Payload))
86 | 		}
87 | 	}()
88 | 
89 | 	<-quit
90 | }
91 | 


--------------------------------------------------------------------------------
/go.mod:
--------------------------------------------------------------------------------
 1 | module github.com/kpeu3i/wspubsub
 2 | 
 3 | go 1.12
 4 | 
 5 | require (
 6 | 	github.com/cespare/xxhash/v2 v2.1.1
 7 | 	github.com/gobwas/httphead v0.0.0-20180130184737-2c6c146eadee // indirect
 8 | 	github.com/gobwas/pool v0.2.0 // indirect
 9 | 	github.com/gobwas/ws v1.0.2
10 | 	github.com/golang/mock v1.4.0
11 | 	github.com/gorilla/websocket v1.4.1
12 | 	github.com/pkg/errors v0.8.1
13 | 	github.com/satori/go.uuid v1.2.0
14 | 	github.com/sirupsen/logrus v1.4.2
15 | 	github.com/stretchr/testify v1.3.0
16 | 	go.uber.org/multierr v1.4.0
17 | 	golang.org/x/sync v0.0.0-20190423024810-112230192c58
18 | )
19 | 


--------------------------------------------------------------------------------
/go.sum:
--------------------------------------------------------------------------------
 1 | github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ=
 2 | github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 3 | github.com/cespare/xxhash/v2 v2.1.1 h1:6MnRN8NT7+YBpUIWxHtefFZOKTAPgGjpQSxqLNn0+qY=
 4 | github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 5 | github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 6 | github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 7 | github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 8 | github.com/gobwas/httphead v0.0.0-20180130184737-2c6c146eadee h1:s+21KNqlpePfkah2I+gwHF8xmJWRjooY+5248k6m4A0=
 9 | github.com/gobwas/httphead v0.0.0-20180130184737-2c6c146eadee/go.mod h1:L0fX3K22YWvt/FAX9NnzrNzcI4wNYi9Yku4O0LKYflo=
10 | github.com/gobwas/pool v0.2.0 h1:QEmUOlnSjWtnpRGHF3SauEiOsy82Cup83Vf2LcMlnc8=
11 | github.com/gobwas/pool v0.2.0/go.mod h1:q8bcK0KcYlCgd9e7WYLm9LpyS+YeLd8JVDW6WezmKEw=
12 | github.com/gobwas/ws v1.0.2 h1:CoAavW/wd/kulfZmSIBt6p24n4j7tHgNVCjsfHVNUbo=
13 | github.com/gobwas/ws v1.0.2/go.mod h1:szmBTxLgaFppYjEmNtny/v3w89xOydFnnZMcgRRu/EM=
14 | github.com/golang/mock v1.4.0 h1:Rd1kQnQu0Hq3qvJppYSG0HtP+f5LPPUiDswTLiEegLg=
15 | github.com/golang/mock v1.4.0/go.mod h1:UOMv5ysSaYNkG+OFQykRIcU/QvvxJf3p21QfJ2Bt3cw=
16 | github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI=
17 | github.com/gorilla/websocket v1.4.1 h1:q7AeDBpnBk8AogcD4DSag/Ukw/KV+YhzLj2bP5HvKCM=
18 | github.com/gorilla/websocket v1.4.1/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
19 | github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
20 | github.com/konsorten/go-windows-terminal-sequences v1.0.1 h1:mweAR1A6xJ3oS2pRaGiHgQ4OO8tzTaLawm8vnODuwDk=
21 | github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
22 | github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
23 | github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
24 | github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
25 | github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
26 | github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
27 | github.com/pkg/errors v0.8.1 h1:iURUrRGxPUNPdy5/HRSm+Yj6okJ6UtLINN0Q9M4+h3I=
28 | github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
29 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
30 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
31 | github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4=
32 | github.com/satori/go.uuid v1.2.0 h1:0uYX9dsZ2yD7q2RtLRtPSdGDWzjeM3TbMJP9utgA0ww=
33 | github.com/satori/go.uuid v1.2.0/go.mod h1:dA0hQrYB0VpLJoorglMZABFdXlWrHn1NEOzdhQKdks0=
34 | github.com/sirupsen/logrus v1.4.2 h1:SPIRibHv4MatM3XXNO2BJeFLZwZ2LvZgfQ5+UNI2im4=
35 | github.com/sirupsen/logrus v1.4.2/go.mod h1:tLMulIdttU9McNUspp0xgXVQah82FyeX6MwdIuYE2rE=
36 | github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
37 | github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
38 | github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
39 | github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0Q=
40 | github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
41 | go.uber.org/atomic v1.5.0 h1:OI5t8sDa1Or+q8AeE+yKeB/SDYioSHAgcVljj9JIETY=
42 | go.uber.org/atomic v1.5.0/go.mod h1:sABNBOSYdrvTF6hTgEIbc7YasKWGhgEQZyfxyTvoXHQ=
43 | go.uber.org/multierr v1.4.0 h1:f3WCSC2KzAcBXGATIxAB1E2XuCpNU255wNKZ505qi3E=
44 | go.uber.org/multierr v1.4.0/go.mod h1:VgVr7evmIr6uPjLBxg28wmKNXyqE9akIJ5XnfpiKl+4=
45 | go.uber.org/tools v0.0.0-20190618225709-2cfd321de3ee h1:0mgffUl7nfd+FpvXMVz4IDEaUSmT1ysygQC7qYo7sG4=
46 | go.uber.org/tools v0.0.0-20190618225709-2cfd321de3ee/go.mod h1:vJERXedbb3MVM5f9Ejo0C68/HhF8uaILCdgjnY+goOA=
47 | golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
48 | golang.org/x/crypto v0.0.0-20190510104115-cbcb75029529/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
49 | golang.org/x/lint v0.0.0-20190930215403-16217165b5de h1:5hukYrvBGR8/eNkX5mdUezrA6JiaEZDtJb9Ei+1LlBs=
50 | golang.org/x/lint v0.0.0-20190930215403-16217165b5de/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
51 | golang.org/x/mod v0.0.0-20190513183733-4bf6d317e70e/go.mod h1:mXi4GBBbnImb6dmsKGUJ2LatrhH/nqhxcFungHvyanc=
52 | golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
53 | golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
54 | golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
55 | golang.org/x/sync v0.0.0-20190423024810-112230192c58 h1:8gQV6CLnAEikrhgkHFbMAEhagSSnXWGV915qUMm9mrU=
56 | golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
57 | golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
58 | golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
59 | golang.org/x/sys v0.0.0-20190422165155-953cdadca894 h1:Cz4ceDQGXuKRnVBDTS23GTn/pU5OE2C0WrNTOYK1Uuc=
60 | golang.org/x/sys v0.0.0-20190422165155-953cdadca894/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
61 | golang.org/x/text v0.0.0-20170915032832-14c0d48ead0c/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
62 | golang.org/x/text v0.3.0 h1:g61tztE5qeGQ89tm6NTjjM9VPIm088od1l6aSorWRWg=
63 | golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
64 | golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
65 | golang.org/x/tools v0.0.0-20190425150028-36563e24a262/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q=
66 | golang.org/x/tools v0.0.0-20190621195816-6e04913cbbac/go.mod h1:/rFqwRUd4F7ZHNgwSSTFct+R/Kf4OFW1sUzUTQQTgfc=
67 | golang.org/x/tools v0.0.0-20191029041327-9cc4af7d6b2c/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
68 | golang.org/x/tools v0.0.0-20191029190741-b9c20aec41a5 h1:hKsoRgsbwY1NafxrwTs+k64bikrLBkAgPir1TNCj3Zs=
69 | golang.org/x/tools v0.0.0-20191029190741-b9c20aec41a5/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
70 | golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
71 | gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
72 | gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
73 | gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
74 | honnef.co/go/tools v0.0.1-2019.2.3 h1:3JgtbtFHMiCmsznwGVTUWbgGov+pVqnlf1dEJTNAXeM=
75 | honnef.co/go/tools v0.0.1-2019.2.3/go.mod h1:a3bituU0lyd329TUQxRnasdCoJDkEUEAqEt0JzvZhAg=
76 | rsc.io/quote/v3 v3.1.0/go.mod h1:yEA65RcK8LyAZtP9Kv3t0HmxON59tX3rD+tICJqUlj0=
77 | rsc.io/sampler v1.3.0/go.mod h1:T1hPZKmBbMNahiBKFy5HrXp6adAjACjK9JXDnKaTXpA=
78 | 


--------------------------------------------------------------------------------
/gobwas_connection.go:
--------------------------------------------------------------------------------
  1 | package wspubsub
  2 | 
  3 | import (
  4 | 	"io/ioutil"
  5 | 	"net"
  6 | 	"strings"
  7 | 	"time"
  8 | 
  9 | 	"github.com/gobwas/ws"
 10 | 	"github.com/gobwas/ws/wsutil"
 11 | 	"github.com/pkg/errors"
 12 | )
 13 | 
 14 | var _ WebsocketConnection = (*GobwasConnection)(nil)
 15 | 
 16 | // GobwasConnection is an implementation of WebsocketConnection.
 17 | type GobwasConnection struct {
 18 | 	conn               net.Conn
 19 | 	logger             Logger
 20 | 	readTimeout        time.Duration
 21 | 	wrightTimout       time.Duration
 22 | 	IsDebug            bool
 23 | 	DebugFuncTimeLimit time.Duration
 24 | }
 25 | 
 26 | // Read reads a message from WebSocket connection.
 27 | func (c *GobwasConnection) Read() (Message, error) {
 28 | 	opCode, bytes, err := c.doRead()
 29 | 	if err != nil {
 30 | 		return Message{}, errors.WithStack(c.handleError(err))
 31 | 	}
 32 | 
 33 | 	message := Message{
 34 | 		Type:    MessageType(opCode),
 35 | 		Payload: bytes,
 36 | 	}
 37 | 
 38 | 	return message, nil
 39 | }
 40 | 
 41 | // Write writes a message to WebSocket connection.
 42 | func (c *GobwasConnection) Write(message Message) error {
 43 | 	if c.IsDebug {
 44 | 		now := time.Now()
 45 | 		defer func() {
 46 | 			end := time.Since(now)
 47 | 			if end > c.DebugFuncTimeLimit {
 48 | 				c.logger.Warnf("gobwas.connection.write: took=%s", end)
 49 | 			}
 50 | 		}()
 51 | 	}
 52 | 
 53 | 	err := c.conn.SetWriteDeadline(time.Now().Add(c.wrightTimout))
 54 | 	if err != nil {
 55 | 		return errors.WithStack(c.handleError(err))
 56 | 	}
 57 | 
 58 | 	err = wsutil.WriteServerMessage(c.conn, ws.OpCode(message.Type), message.Payload)
 59 | 	if err != nil {
 60 | 		return errors.WithStack(c.handleError(err))
 61 | 	}
 62 | 
 63 | 	return nil
 64 | }
 65 | 
 66 | // Close closes a WebSocket connection.
 67 | func (c *GobwasConnection) Close() error {
 68 | 	if c.IsDebug {
 69 | 		now := time.Now()
 70 | 		defer func() {
 71 | 			end := time.Since(now)
 72 | 			if end > c.DebugFuncTimeLimit {
 73 | 				c.logger.Warnf("gobwas.connection.close: took=%s", end)
 74 | 			}
 75 | 		}()
 76 | 	}
 77 | 
 78 | 	err := c.conn.Close()
 79 | 	if err != nil {
 80 | 		return errors.WithStack(c.handleError(err))
 81 | 	}
 82 | 
 83 | 	return nil
 84 | }
 85 | 
 86 | func (c *GobwasConnection) doRead() (ws.OpCode, []byte, error) {
 87 | 	controlHandler := wsutil.ControlFrameHandler(c.conn, ws.StateServerSide)
 88 | 	reader := wsutil.Reader{
 89 | 		Source:          c.conn,
 90 | 		State:           ws.StateServerSide,
 91 | 		CheckUTF8:       true,
 92 | 		SkipHeaderCheck: false,
 93 | 		OnIntermediate:  controlHandler,
 94 | 	}
 95 | 
 96 | 	err := c.conn.SetReadDeadline(time.Now().Add(c.readTimeout))
 97 | 	if err != nil {
 98 | 		return 0, nil, err
 99 | 	}
100 | 
101 | 	for {
102 | 		header, err := reader.NextFrame()
103 | 		if err != nil {
104 | 			return 0, nil, err
105 | 		}
106 | 
107 | 		if header.OpCode.IsControl() {
108 | 			if header.OpCode == ws.OpPong {
109 | 				err := c.conn.SetReadDeadline(time.Now().Add(c.readTimeout))
110 | 				if err != nil {
111 | 					return 0, nil, err
112 | 				}
113 | 			}
114 | 
115 | 			err := controlHandler(header, &reader)
116 | 			if err != nil {
117 | 				return 0, nil, err
118 | 			}
119 | 
120 | 			continue
121 | 		}
122 | 
123 | 		if header.OpCode&(ws.OpText|ws.OpBinary) == 0 {
124 | 			err := reader.Discard()
125 | 			if err != nil {
126 | 				return 0, nil, err
127 | 			}
128 | 
129 | 			continue
130 | 		}
131 | 
132 | 		bytes, err := ioutil.ReadAll(&reader)
133 | 
134 | 		return header.OpCode, bytes, err
135 | 	}
136 | }
137 | 
138 | func (c *GobwasConnection) handleError(err error) error {
139 | 	if err == nil {
140 | 		return nil
141 | 	}
142 | 
143 | 	if _, ok := err.(wsutil.ClosedError); ok {
144 | 		closeErr := NewConnectionClosedError(err)
145 | 
146 | 		return errors.WithStack(closeErr)
147 | 	}
148 | 
149 | 	if strings.Contains(err.Error(), "use of closed network connection") {
150 | 		closeErr := NewConnectionClosedError(err)
151 | 
152 | 		return errors.WithStack(closeErr)
153 | 	}
154 | 
155 | 	if _, ok := err.(*net.OpError); ok {
156 | 		closeErr := NewConnectionClosedError(err)
157 | 
158 | 		return errors.WithStack(closeErr)
159 | 	}
160 | 
161 | 	return errors.WithStack(err)
162 | }
163 | 


--------------------------------------------------------------------------------
/gobwas_upgrader.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import (
 4 | 	"net/http"
 5 | 	"time"
 6 | 
 7 | 	"github.com/gobwas/ws"
 8 | )
 9 | 
10 | var _ WebsocketConnectionUpgrader = (*GobwasConnectionUpgrader)(nil)
11 | 
12 | // GobwasConnectionUpgrader is an implementation of WebsocketConnectionUpgrader.
13 | type GobwasConnectionUpgrader struct {
14 | 	logger  Logger
15 | 	options GobwasConnectionUpgraderOptions
16 | }
17 | 
18 | // GobwasConnectionUpgrader upgrades HTTP connection to the WebSocket connection.
19 | func (u *GobwasConnectionUpgrader) Upgrade(w http.ResponseWriter, r *http.Request) (WebsocketConnection, error) {
20 | 	if u.options.IsDebug {
21 | 		now := time.Now()
22 | 		defer func() {
23 | 			end := time.Since(now)
24 | 			if end > u.options.DebugFuncTimeLimit {
25 | 				u.logger.Warnf("gobwas.connection_upgrader.upgrader: took=%s", end)
26 | 			}
27 | 		}()
28 | 	}
29 | 
30 | 	connection, _, _, err := ws.UpgradeHTTP(r, w)
31 | 	if err != nil {
32 | 		return nil, err
33 | 	}
34 | 
35 | 	gobwasConnection := &GobwasConnection{
36 | 		conn:               connection,
37 | 		logger:             u.logger,
38 | 		readTimeout:        u.options.ReadTimout,
39 | 		wrightTimout:       u.options.WriteTimout,
40 | 		IsDebug:            u.options.IsDebug,
41 | 		DebugFuncTimeLimit: u.options.DebugFuncTimeLimit,
42 | 	}
43 | 
44 | 	return gobwasConnection, nil
45 | }
46 | 
47 | // NewGobwasConnectionUpgrader initializes a new GobwasConnectionUpgrader.
48 | func NewGobwasConnectionUpgrader(options GobwasConnectionUpgraderOptions, logger Logger) *GobwasConnectionUpgrader {
49 | 	return &GobwasConnectionUpgrader{options: options, logger: logger}
50 | }
51 | 


--------------------------------------------------------------------------------
/gobwas_upgrader_options.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import (
 4 | 	"time"
 5 | )
 6 | 
 7 | // GobwasConnectionUpgraderOptions represents configuration of the GobwasConnectionUpgrader.
 8 | type GobwasConnectionUpgraderOptions struct {
 9 | 	ReadTimout         time.Duration
10 | 	WriteTimout        time.Duration
11 | 	IsDebug            bool
12 | 	DebugFuncTimeLimit time.Duration
13 | }
14 | 
15 | // NewGobwasUpgraderOptions initializes a new GobwasConnectionUpgraderOptions.
16 | // nolint: gomnd
17 | func NewGobwasConnectionUpgraderOptions() GobwasConnectionUpgraderOptions {
18 | 	options := GobwasConnectionUpgraderOptions{
19 | 		ReadTimout:         60 * time.Second,
20 | 		WriteTimout:        10 * time.Second,
21 | 		IsDebug:            false,
22 | 		DebugFuncTimeLimit: 1 * time.Millisecond,
23 | 	}
24 | 
25 | 	return options
26 | }
27 | 


--------------------------------------------------------------------------------
/gobwas_upgrader_options_test.go:
--------------------------------------------------------------------------------
 1 | package wspubsub_test
 2 | 
 3 | import (
 4 | 	"testing"
 5 | 
 6 | 	"github.com/kpeu3i/wspubsub"
 7 | 	"github.com/stretchr/testify/require"
 8 | )
 9 | 
10 | func TestNewGobwasUpgraderOptions(t *testing.T) {
11 | 	options := wspubsub.NewGobwasConnectionUpgraderOptions()
12 | 	require.NotZero(t, options.ReadTimout)
13 | 	require.NotZero(t, options.WriteTimout)
14 | 	require.False(t, options.IsDebug)
15 | 	require.NotZero(t, options.DebugFuncTimeLimit)
16 | }
17 | 


--------------------------------------------------------------------------------
/gorilla_connection.go:
--------------------------------------------------------------------------------
  1 | package wspubsub
  2 | 
  3 | import (
  4 | 	"net"
  5 | 	"strings"
  6 | 	"time"
  7 | 
  8 | 	"github.com/gorilla/websocket"
  9 | 	"github.com/pkg/errors"
 10 | )
 11 | 
 12 | var _ WebsocketConnection = (*GorillaConnection)(nil)
 13 | 
 14 | // GorillaConnection is an implementation of WebsocketConnection.
 15 | type GorillaConnection struct {
 16 | 	conn               *websocket.Conn
 17 | 	logger             Logger
 18 | 	maxMessageSize     int64
 19 | 	readTimeout        time.Duration
 20 | 	writeTimout        time.Duration
 21 | 	IsDebug            bool
 22 | 	DebugFuncTimeLimit time.Duration
 23 | }
 24 | 
 25 | // Read reads a message from WebSocket connection.
 26 | func (c *GorillaConnection) Read() (Message, error) {
 27 | 	err := c.conn.SetReadDeadline(time.Now().Add(c.readTimeout))
 28 | 	if err != nil {
 29 | 		return Message{}, errors.WithStack(c.handleError(err))
 30 | 	}
 31 | 
 32 | 	messageType, bytes, err := c.conn.ReadMessage()
 33 | 	if err != nil {
 34 | 		return Message{}, errors.WithStack(c.handleError(err))
 35 | 	}
 36 | 
 37 | 	message := Message{
 38 | 		Type:    MessageType(messageType),
 39 | 		Payload: bytes,
 40 | 	}
 41 | 
 42 | 	return message, nil
 43 | }
 44 | 
 45 | // Write writes a message to WebSocket connection.
 46 | func (c *GorillaConnection) Write(message Message) error {
 47 | 	if c.IsDebug {
 48 | 		now := time.Now()
 49 | 		defer func() {
 50 | 			end := time.Since(now)
 51 | 			if end > c.DebugFuncTimeLimit {
 52 | 				c.logger.Warnf("wspubsub.gorilla_connection.write: took=%s", end)
 53 | 			}
 54 | 		}()
 55 | 	}
 56 | 
 57 | 	err := c.conn.SetWriteDeadline(time.Now().Add(c.writeTimout))
 58 | 	if err != nil {
 59 | 		return errors.WithStack(c.handleError(err))
 60 | 	}
 61 | 
 62 | 	err = c.conn.WriteMessage(int(message.Type), message.Payload)
 63 | 	if err != nil {
 64 | 		return errors.WithStack(c.handleError(err))
 65 | 	}
 66 | 
 67 | 	return nil
 68 | }
 69 | 
 70 | // Close closes a WebSocket connection.
 71 | func (c *GorillaConnection) Close() error {
 72 | 	if c.IsDebug {
 73 | 		now := time.Now()
 74 | 		defer func() {
 75 | 			end := time.Since(now)
 76 | 			if end > c.DebugFuncTimeLimit {
 77 | 				c.logger.Warnf("wspubsub.gorilla_connection.close: took=%s", end)
 78 | 			}
 79 | 		}()
 80 | 	}
 81 | 
 82 | 	err := c.conn.Close()
 83 | 	if err != nil {
 84 | 		return errors.WithStack(c.handleError(err))
 85 | 	}
 86 | 
 87 | 	return nil
 88 | }
 89 | 
 90 | func (c *GorillaConnection) handleError(err error) error {
 91 | 	if err == nil {
 92 | 		return nil
 93 | 	}
 94 | 
 95 | 	if _, ok := err.(*websocket.CloseError); ok {
 96 | 		closeErr := NewConnectionClosedError(err)
 97 | 
 98 | 		return errors.WithStack(closeErr)
 99 | 	}
100 | 
101 | 	if err == websocket.ErrCloseSent {
102 | 		closeErr := NewConnectionClosedError(err)
103 | 
104 | 		return errors.WithStack(closeErr)
105 | 	}
106 | 
107 | 	if strings.Contains(err.Error(), "use of closed network connection") {
108 | 		closeErr := NewConnectionClosedError(err)
109 | 
110 | 		return errors.WithStack(closeErr)
111 | 	}
112 | 
113 | 	if _, ok := err.(*net.OpError); ok {
114 | 		closeErr := NewConnectionClosedError(err)
115 | 
116 | 		return errors.WithStack(closeErr)
117 | 	}
118 | 
119 | 	return errors.WithStack(err)
120 | }
121 | 


--------------------------------------------------------------------------------
/gorilla_upgrader.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import (
 4 | 	"net/http"
 5 | 	"time"
 6 | 
 7 | 	"github.com/gorilla/websocket"
 8 | 	"github.com/pkg/errors"
 9 | )
10 | 
11 | var _ WebsocketConnectionUpgrader = (*GorillaConnectionUpgrader)(nil)
12 | 
13 | // GorillaConnectionUpgrader is an implementation of WebsocketConnectionUpgrader.
14 | type GorillaConnectionUpgrader struct {
15 | 	options  GorillaConnectionUpgraderOptions
16 | 	logger   Logger
17 | 	upgrader *websocket.Upgrader
18 | }
19 | 
20 | // GorillaConnectionUpgrader upgrades HTTP connection to the WebSocket connection.
21 | func (u *GorillaConnectionUpgrader) Upgrade(w http.ResponseWriter, r *http.Request) (WebsocketConnection, error) {
22 | 	if u.options.IsDebug {
23 | 		now := time.Now()
24 | 		defer func() {
25 | 			end := time.Since(now)
26 | 			if end > u.options.DebugFuncTimeLimit {
27 | 				u.logger.Warnf("gorilla.connection_upgrader.upgrader: took=%s", end)
28 | 			}
29 | 		}()
30 | 	}
31 | 
32 | 	connection, err := u.upgrader.Upgrade(w, r, nil)
33 | 	if err != nil {
34 | 		return nil, errors.WithStack(err)
35 | 	}
36 | 
37 | 	connection.SetReadLimit(u.options.MaxMessageSize)
38 | 
39 | 	err = connection.SetReadDeadline(time.Now().Add(u.options.ReadTimout))
40 | 	if err != nil {
41 | 		return nil, errors.WithStack(err)
42 | 	}
43 | 
44 | 	connection.SetPongHandler(func(string) error {
45 | 		return connection.SetReadDeadline(time.Now().Add(u.options.ReadTimout))
46 | 	})
47 | 
48 | 	gorillaConnection := &GorillaConnection{
49 | 		conn:               connection,
50 | 		logger:             u.logger,
51 | 		maxMessageSize:     u.options.MaxMessageSize,
52 | 		readTimeout:        u.options.ReadTimout,
53 | 		writeTimout:        u.options.WriteTimout,
54 | 		IsDebug:            u.options.IsDebug,
55 | 		DebugFuncTimeLimit: u.options.DebugFuncTimeLimit,
56 | 	}
57 | 
58 | 	return gorillaConnection, nil
59 | }
60 | 
61 | // NewGorillaConnectionUpgrader initializes a new GorillaConnectionUpgrader.
62 | func NewGorillaConnectionUpgrader(options GorillaConnectionUpgraderOptions, logger Logger) *GorillaConnectionUpgrader {
63 | 	upgrader := &websocket.Upgrader{
64 | 		HandshakeTimeout:  options.HandshakeTimeout,
65 | 		ReadBufferSize:    options.ReadBufferSize,
66 | 		WriteBufferSize:   options.WriteBufferSize,
67 | 		Subprotocols:      options.Subprotocols,
68 | 		CheckOrigin:       options.CheckOrigin,
69 | 		EnableCompression: options.EnableCompression,
70 | 	}
71 | 
72 | 	return &GorillaConnectionUpgrader{options: options, logger: logger, upgrader: upgrader}
73 | }
74 | 


--------------------------------------------------------------------------------
/gorilla_upgrader_options.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import (
 4 | 	"net/http"
 5 | 	"time"
 6 | )
 7 | 
 8 | // GorillaConnectionUpgraderOptions represents configuration of the GorillaConnectionUpgrader.
 9 | type GorillaConnectionUpgraderOptions struct {
10 | 	MaxMessageSize     int64
11 | 	ReadTimout         time.Duration
12 | 	WriteTimout        time.Duration
13 | 	HandshakeTimeout   time.Duration
14 | 	ReadBufferSize     int
15 | 	WriteBufferSize    int
16 | 	Subprotocols       []string
17 | 	Error              func(w http.ResponseWriter, r *http.Request, status int, reason error)
18 | 	CheckOrigin        func(r *http.Request) bool
19 | 	EnableCompression  bool
20 | 	IsDebug            bool
21 | 	DebugFuncTimeLimit time.Duration
22 | }
23 | 
24 | // NewGorillaConnectionUpgraderOptions initializes a new GorillaConnectionUpgraderOptions.
25 | // nolint: gomnd
26 | func NewGorillaConnectionUpgraderOptions() GorillaConnectionUpgraderOptions {
27 | 	options := GorillaConnectionUpgraderOptions{
28 | 		MaxMessageSize:  1 * 1024 * 1024,
29 | 		ReadTimout:      60 * time.Second,
30 | 		WriteTimout:     10 * time.Second,
31 | 		ReadBufferSize:  4096,
32 | 		WriteBufferSize: 4096,
33 | 		CheckOrigin: func(r *http.Request) bool {
34 | 			return true
35 | 		},
36 | 		IsDebug:            false,
37 | 		DebugFuncTimeLimit: 1 * time.Millisecond,
38 | 	}
39 | 
40 | 	return options
41 | }
42 | 


--------------------------------------------------------------------------------
/gorilla_upgrader_options_test.go:
--------------------------------------------------------------------------------
 1 | package wspubsub_test
 2 | 
 3 | import (
 4 | 	"testing"
 5 | 
 6 | 	"github.com/kpeu3i/wspubsub"
 7 | 	"github.com/stretchr/testify/require"
 8 | )
 9 | 
10 | func TestNewGorillaUpgraderOptions(t *testing.T) {
11 | 	options := wspubsub.NewGorillaConnectionUpgraderOptions()
12 | 	require.NotZero(t, options.MaxMessageSize)
13 | 	require.NotZero(t, options.ReadTimout)
14 | 	require.NotZero(t, options.WriteTimout)
15 | 	require.False(t, options.IsDebug)
16 | 	require.NotZero(t, options.DebugFuncTimeLimit)
17 | }
18 | 


--------------------------------------------------------------------------------
/hub.go:
--------------------------------------------------------------------------------
  1 | package wspubsub
  2 | 
  3 | import (
  4 | 	"context"
  5 | 	"net/http"
  6 | 	"strings"
  7 | 	"sync/atomic"
  8 | 	"time"
  9 | 
 10 | 	"github.com/pkg/errors"
 11 | 	"go.uber.org/multierr"
 12 | 	"golang.org/x/sync/errgroup"
 13 | )
 14 | 
 15 | // WebsocketClient is an interface representing the ability to interact with WebSocket connection.
 16 | type WebsocketClient interface {
 17 | 	ID() UUID
 18 | 	Connect(response http.ResponseWriter, request *http.Request) error
 19 | 	OnReceive(handler ReceiveHandler)
 20 | 	OnError(handler ErrorHandler)
 21 | 	Send(message Message) error
 22 | 	Close() error
 23 | }
 24 | 
 25 | // WebsocketClientStore is an interface responsible for storing and finding the users.
 26 | type WebsocketClientStore interface {
 27 | 	Get(clientID UUID) (WebsocketClient, error)
 28 | 	Set(client WebsocketClient)
 29 | 	Unset(clientID UUID) error
 30 | 	Count(channels ...string) int
 31 | 	Find(fn IterateFunc, channels ...string) error
 32 | 	Channels(clientID UUID) ([]string, error)
 33 | 	CountChannels(clientID UUID) (int, error)
 34 | 	SetChannels(clientID UUID, channels ...string) error
 35 | 	UnsetChannels(clientID UUID, channels ...string) error
 36 | }
 37 | 
 38 | // WebsocketClientStore is an interface responsible for creating a client.
 39 | type WebsocketClientFactory interface {
 40 | 	Create() WebsocketClient
 41 | }
 42 | 
 43 | // Logger is an interface representing the ability to log messages.
 44 | type Logger interface {
 45 | 	Debug(args ...interface{})
 46 | 	Info(args ...interface{})
 47 | 	Print(args ...interface{})
 48 | 	Warn(args ...interface{})
 49 | 	Error(args ...interface{})
 50 | 	Fatal(args ...interface{})
 51 | 	Panic(args ...interface{})
 52 | 
 53 | 	Debugln(args ...interface{})
 54 | 	Infoln(args ...interface{})
 55 | 	Println(args ...interface{})
 56 | 	Warnln(args ...interface{})
 57 | 	Errorln(args ...interface{})
 58 | 	Fatalln(args ...interface{})
 59 | 	Panicln(args ...interface{})
 60 | 
 61 | 	Debugf(format string, args ...interface{})
 62 | 	Infof(format string, args ...interface{})
 63 | 	Printf(format string, args ...interface{})
 64 | 	Warnf(format string, args ...interface{})
 65 | 	Errorf(format string, args ...interface{})
 66 | 	Fatalf(format string, args ...interface{})
 67 | 	Panicf(format string, args ...interface{})
 68 | }
 69 | 
 70 | type (
 71 | 	// ConnectHandler called when a new client is connected to hub.
 72 | 	ConnectHandler func(clientID UUID)
 73 | 
 74 | 	// DisconnectHandler called when a client is disconnected from the hub.
 75 | 	DisconnectHandler func(clientID UUID)
 76 | 
 77 | 	// ReceiveHandler called when a client reads a new message.
 78 | 	ReceiveHandler func(clientID UUID, message Message)
 79 | 
 80 | 	// ErrorHandler called when an error occurred when reading or writing messages.
 81 | 	ErrorHandler func(clientID UUID, err error)
 82 | )
 83 | 
 84 | // nolint: gochecknoglobals
 85 | var (
 86 | 	defaultConnectHandler    = ConnectHandler(func(clientID UUID) {})
 87 | 	defaultDisconnectHandler = DisconnectHandler(func(clientID UUID) {})
 88 | 	defaultReceiveHandler    = ReceiveHandler(func(clientID UUID, message Message) {})
 89 | 	defaultErrorHandler      = ErrorHandler(func(clientID UUID, err error) {})
 90 | )
 91 | 
 92 | // Hub manages client connections.
 93 | type Hub struct {
 94 | 	options           HubOptions
 95 | 	clients           WebsocketClientStore
 96 | 	clientFactory     WebsocketClientFactory
 97 | 	logger            Logger
 98 | 	httpServer        *http.Server
 99 | 	httpServerTLS     *http.Server
100 | 	connectHandler    atomic.Value
101 | 	disconnectHandler atomic.Value
102 | 	receiveHandler    atomic.Value
103 | 	errorHandler      atomic.Value
104 | }
105 | 
106 | // Subscribe allows to subscribe a client to specific channels.
107 | // At least one channel is required.
108 | func (h *Hub) Subscribe(clientID UUID, channels ...string) error {
109 | 	if h.options.IsDebug {
110 | 		now := time.Now()
111 | 		defer func() {
112 | 			end := time.Since(now)
113 | 			if end > h.options.DebugFuncTimeLimit {
114 | 				h.logger.Warnf("wspubsub.hub.subscribe: took=%s", end)
115 | 			}
116 | 		}()
117 | 	}
118 | 
119 | 	if len(channels) == 0 {
120 | 		return NewHubSubscriptionChannelRequiredError()
121 | 	}
122 | 
123 | 	err := h.clients.SetChannels(clientID, channels...)
124 | 	if err != nil {
125 | 		return errors.WithStack(err)
126 | 	}
127 | 
128 | 	if h.options.IsDebug {
129 | 		h.logger.Debugf("Client subscribed: id=%s, channels=[%s]", clientID, strings.Join(channels, ","))
130 | 	}
131 | 
132 | 	return nil
133 | }
134 | 
135 | // Unsubscribe allows to unsubscribe a client from specific channels.
136 | // If channels were not specified then the client will be
137 | // unsubscribed from all channels.
138 | func (h *Hub) Unsubscribe(clientID UUID, channels ...string) error {
139 | 	if h.options.IsDebug {
140 | 		now := time.Now()
141 | 		defer func() {
142 | 			end := time.Since(now)
143 | 			if end > h.options.DebugFuncTimeLimit {
144 | 				h.logger.Warnf("wspubsub.hub.unsubscribe: took=%s", end)
145 | 			}
146 | 		}()
147 | 	}
148 | 
149 | 	err := h.clients.UnsetChannels(clientID, channels...)
150 | 	if err != nil {
151 | 		return errors.WithStack(err)
152 | 	}
153 | 
154 | 	if h.options.IsDebug {
155 | 		h.logger.Debugf("Client unsubscribed: id=%s, channels=[%s]", clientID, strings.Join(channels, ","))
156 | 	}
157 | 
158 | 	return nil
159 | }
160 | 
161 | // IsSubscribed checks does the client is subscribed to at least one channel.
162 | func (h *Hub) IsSubscribed(clientID UUID) bool {
163 | 	if h.options.IsDebug {
164 | 		now := time.Now()
165 | 		defer func() {
166 | 			end := time.Since(now)
167 | 			if end > h.options.DebugFuncTimeLimit {
168 | 				h.logger.Warnf("wspubsub.hub.is_subscribed: took=%s", end)
169 | 			}
170 | 		}()
171 | 	}
172 | 
173 | 	count, _ := h.clients.CountChannels(clientID)
174 | 
175 | 	return count > 0
176 | }
177 | 
178 | // Channels return a list of channels the client currently subscribed to.
179 | func (h *Hub) Channels(clientID UUID) ([]string, error) {
180 | 	if h.options.IsDebug {
181 | 		now := time.Now()
182 | 		defer func() {
183 | 			end := time.Since(now)
184 | 			if end > h.options.DebugFuncTimeLimit {
185 | 				h.logger.Warnf("wspubsub.hub.channels: took=%s", end)
186 | 			}
187 | 		}()
188 | 	}
189 | 
190 | 	channels, err := h.clients.Channels(clientID)
191 | 	if err != nil {
192 | 		return nil, errors.WithStack(err)
193 | 	}
194 | 
195 | 	return channels, nil
196 | }
197 | 
198 | // Count returns total number of connected clients.
199 | func (h *Hub) Count(channels ...string) int {
200 | 	if h.options.IsDebug {
201 | 		now := time.Now()
202 | 		defer func() {
203 | 			end := time.Since(now)
204 | 			if end > h.options.DebugFuncTimeLimit {
205 | 				h.logger.Warnf("wspubsub.hub.count: took=%s", end)
206 | 			}
207 | 		}()
208 | 	}
209 | 
210 | 	return h.clients.Count(channels...)
211 | }
212 | 
213 | // Publish publishes a message to the channels.
214 | // If channels were not specified then all clients will receive the message.
215 | func (h *Hub) Publish(message Message, channels ...string) (int, error) {
216 | 	if h.options.IsDebug {
217 | 		now := time.Now()
218 | 		defer func() {
219 | 			end := time.Since(now)
220 | 			if end > h.options.DebugFuncTimeLimit {
221 | 				h.logger.Warnf("wspubsub.hub.publish: took=%s", end)
222 | 			}
223 | 		}()
224 | 	}
225 | 
226 | 	numClients := 0
227 | 	iterateFunc := func(client WebsocketClient) error {
228 | 		err := client.Send(message)
229 | 		if err != nil {
230 | 			// A buffer overflow error can occur here,
231 | 			// so we should disconnect the client
232 | 			_ = h.disconnectClient(client)
233 | 
234 | 			return nil
235 | 		}
236 | 
237 | 		numClients++
238 | 
239 | 		return nil
240 | 	}
241 | 
242 | 	err := h.clients.Find(iterateFunc, channels...)
243 | 	if err != nil {
244 | 		return numClients, errors.WithStack(err)
245 | 	}
246 | 
247 | 	if h.options.IsDebug {
248 | 		if numClients > 0 {
249 | 			h.logger.Debugf("Message published: num_clients=%d, channels=[%s]", numClients, strings.Join(channels, ","))
250 | 		}
251 | 	}
252 | 
253 | 	return numClients, nil
254 | }
255 | 
256 | // Send sends a message to a specific client.
257 | func (h *Hub) Send(clientID UUID, message Message) error {
258 | 	if h.options.IsDebug {
259 | 		now := time.Now()
260 | 		defer func() {
261 | 			end := time.Since(now)
262 | 			if end > h.options.DebugFuncTimeLimit {
263 | 				h.logger.Warnf("wspubsub.hub.send: took=%s", end)
264 | 			}
265 | 		}()
266 | 	}
267 | 
268 | 	client, err := h.clients.Get(clientID)
269 | 	if err != nil {
270 | 		return errors.WithStack(err)
271 | 	}
272 | 
273 | 	err = client.Send(message)
274 | 	if err != nil {
275 | 		return errors.WithStack(err)
276 | 	}
277 | 
278 | 	if h.options.IsDebug {
279 | 		h.logger.Debugf("Message sent: id=%s", clientID)
280 | 	}
281 | 
282 | 	return nil
283 | }
284 | 
285 | // Disconnect closes a client connection and removes it from the storage.
286 | func (h *Hub) Disconnect(clientID UUID) error {
287 | 	if h.options.IsDebug {
288 | 		now := time.Now()
289 | 		defer func() {
290 | 			end := time.Since(now)
291 | 			if end > h.options.DebugFuncTimeLimit {
292 | 				h.logger.Warnf("wspubsub.hub.disconnect: took=%s", end)
293 | 			}
294 | 		}()
295 | 	}
296 | 
297 | 	client, err := h.clients.Get(clientID)
298 | 	if err != nil {
299 | 		return errors.WithStack(err)
300 | 	}
301 | 
302 | 	err = h.disconnectClient(client)
303 | 	if err != nil {
304 | 		return errors.WithStack(err)
305 | 	}
306 | 
307 | 	if h.options.IsDebug {
308 | 		h.logger.Debugf("Client disconnected: id=%s", clientID)
309 | 	}
310 | 
311 | 	return nil
312 | }
313 | 
314 | // ListenAndServe listens on the TCP network address and handle requests
315 | // on incoming connections.
316 | func (h *Hub) ListenAndServe(addr, path string) error {
317 | 	h.logger.Infof("Listening connection on: addr=%s, path=%s", addr, path)
318 | 
319 | 	mux := http.NewServeMux()
320 | 	mux.Handle(path, h)
321 | 
322 | 	h.httpServer.Addr = addr
323 | 	h.httpServer.Handler = mux
324 | 
325 | 	err := h.httpServer.ListenAndServe()
326 | 	if err != nil {
327 | 		if err != http.ErrServerClosed {
328 | 			return errors.WithStack(err)
329 | 		}
330 | 	}
331 | 
332 | 	return nil
333 | }
334 | 
335 | // ListenAndServe listens on the TCP network address and handle requests
336 | // on incoming connections.
337 | func (h *Hub) ListenAndServeTLS(addr, path, certFile, keyFile string) error {
338 | 	h.logger.Infof("Listening TLS connection on: addr=%s, path=%s", addr, path)
339 | 
340 | 	mux := http.NewServeMux()
341 | 	mux.Handle(path, h)
342 | 
343 | 	h.httpServerTLS.Addr = addr
344 | 	h.httpServerTLS.Handler = mux
345 | 
346 | 	err := h.httpServerTLS.ListenAndServeTLS(certFile, keyFile)
347 | 	if err != nil {
348 | 		if err != http.ErrServerClosed {
349 | 			return errors.WithStack(err)
350 | 		}
351 | 	}
352 | 
353 | 	return nil
354 | }
355 | 
356 | // ServeHTTP implements http.Handler interface and responsible for connect new clients.
357 | func (h *Hub) ServeHTTP(response http.ResponseWriter, request *http.Request) {
358 | 	if h.options.IsDebug {
359 | 		now := time.Now()
360 | 		defer func() {
361 | 			end := time.Since(now)
362 | 			if end > h.options.DebugFuncTimeLimit {
363 | 				h.logger.Warnf("wspubsub.hub.connection_upgrade_handler: took=%s", end)
364 | 			}
365 | 		}()
366 | 	}
367 | 
368 | 	receiveHandler := h.receiveHandler.Load().(ReceiveHandler)
369 | 	errorHandler := h.errorHandler.Load().(ErrorHandler)
370 | 
371 | 	client := h.clientFactory.Create()
372 | 	client.OnReceive(receiveHandler)
373 | 	client.OnError(errorHandler)
374 | 
375 | 	err := h.connectClient(client, response, request)
376 | 	if err != nil {
377 | 		http.Error(response, "Internal Server Error", http.StatusInternalServerError)
378 | 		if h.options.IsDebug {
379 | 			h.logger.Errorf("cant upgrade connection: %s", err)
380 | 		}
381 | 
382 | 		return
383 | 	}
384 | 
385 | 	if h.options.IsDebug {
386 | 		h.logger.Debugf("connection upgraded: id=%s", client.ID())
387 | 	}
388 | }
389 | 
390 | // Close shutdowns http-servers and disconnects clients.
391 | func (h *Hub) Close() error {
392 | 	if h.options.IsDebug {
393 | 		now := time.Now()
394 | 		defer func() {
395 | 			end := time.Since(now)
396 | 			if end > h.options.DebugFuncTimeLimit {
397 | 				h.logger.Warnf("wspubsub.hub.close: took=%s", end)
398 | 			}
399 | 		}()
400 | 	}
401 | 
402 | 	h.logger.Info("Closing connections...")
403 | 
404 | 	ctx, cancel := context.WithTimeout(context.Background(), h.options.ShutdownTimeout)
405 | 	defer cancel()
406 | 
407 | 	eg := errgroup.Group{}
408 | 
409 | 	eg.Go(func() error {
410 | 		return h.httpServer.Shutdown(ctx)
411 | 	})
412 | 
413 | 	eg.Go(func() error {
414 | 		return h.httpServerTLS.Shutdown(ctx)
415 | 	})
416 | 
417 | 	errList := multierr.Combine(eg.Wait())
418 | 
419 | 	iterateFunc := func(client WebsocketClient) error {
420 | 		_ = h.disconnectClient(client)
421 | 
422 | 		return nil
423 | 	}
424 | 
425 | 	errList = multierr.Combine(errList, h.clients.Find(iterateFunc))
426 | 
427 | 	return errors.WithStack(errList)
428 | }
429 | 
430 | // OnConnect registers a handler for client connection.
431 | func (h *Hub) OnConnect(handler ConnectHandler) {
432 | 	h.logger.Infof("Registering handler: %T", handler)
433 | 	h.connectHandler.Store(handler)
434 | }
435 | 
436 | // OnDisconnect registers a handler for client disconnection.
437 | func (h *Hub) OnDisconnect(handler DisconnectHandler) {
438 | 	h.logger.Infof("Registering handler: %T", handler)
439 | 	h.disconnectHandler.Store(handler)
440 | }
441 | 
442 | // OnReceive registers a handler for incoming messages.
443 | func (h *Hub) OnReceive(handler ReceiveHandler) {
444 | 	h.logger.Infof("Registering handler: %T", handler)
445 | 	h.receiveHandler.Store(handler)
446 | }
447 | 
448 | // OnError registers a handler for errors occurred while reading or writing connection.
449 | func (h *Hub) OnError(handler ErrorHandler) {
450 | 	h.logger.Infof("Registering handler: %T", handler)
451 | 	h.errorHandler.Store(h.wrapErrorHandler(handler))
452 | }
453 | 
454 | // LogDebug logs a message at level Debug.
455 | func (h *Hub) LogDebug(args ...interface{}) {
456 | 	h.logger.Debug(args...)
457 | }
458 | 
459 | // LogInfo logs a message at level Info.
460 | func (h *Hub) LogInfo(args ...interface{}) {
461 | 	h.logger.Info(args...)
462 | }
463 | 
464 | // LogPrint logs a message at level Info.
465 | func (h *Hub) LogPrint(args ...interface{}) {
466 | 	h.logger.Print(args...)
467 | }
468 | 
469 | // LogWarn logs a message at level Warn.
470 | func (h *Hub) LogWarn(args ...interface{}) {
471 | 	h.logger.Warn(args...)
472 | }
473 | 
474 | // LogError logs a message at level Error.
475 | func (h *Hub) LogError(args ...interface{}) {
476 | 	h.logger.Error(args...)
477 | }
478 | 
479 | // Fatal logs a message at level Fatal then the process will exit with status set to 1.
480 | func (h *Hub) LogFatal(args ...interface{}) {
481 | 	h.logger.Fatal(args...)
482 | }
483 | 
484 | // Panic logs a message at level Panic and panics.
485 | func (h *Hub) LogPanic(args ...interface{}) {
486 | 	h.logger.Panic(args...)
487 | }
488 | 
489 | // LogDebugln is like LogDebug but adds a new line.
490 | func (h *Hub) LogDebugln(args ...interface{}) {
491 | 	h.logger.Debugln(args...)
492 | }
493 | 
494 | // LogInfoln is like LogInfo but adds a new line.
495 | func (h *Hub) LogInfoln(args ...interface{}) {
496 | 	h.logger.Infoln(args...)
497 | }
498 | 
499 | // LogPrintln is like LogPrint but adds a new line.
500 | func (h *Hub) LogPrintln(args ...interface{}) {
501 | 	h.logger.Println(args...)
502 | }
503 | 
504 | // LogWarnln is like LogWarn but adds a new line.
505 | func (h *Hub) LogWarnln(args ...interface{}) {
506 | 	h.logger.Warnln(args...)
507 | }
508 | 
509 | // LogErrorln is like LogError but adds a new line.
510 | func (h *Hub) LogErrorln(args ...interface{}) {
511 | 	h.logger.Errorln(args...)
512 | }
513 | 
514 | // LogFatalln is like LogFatal but adds a new line.
515 | func (h *Hub) LogFatalln(args ...interface{}) {
516 | 	h.logger.Fatalln(args...)
517 | }
518 | 
519 | // LogPanicln is like LogPanic but adds a new line.
520 | func (h *Hub) LogPanicln(args ...interface{}) {
521 | 	h.logger.Panicln(args...)
522 | }
523 | 
524 | // LogDebugf is like LogDebug but allows specifying a message format.
525 | func (h *Hub) LogDebugf(format string, args ...interface{}) {
526 | 	h.logger.Debugf(format, args...)
527 | }
528 | 
529 | // LogInfof is like LogInfo but allows specifying a message format.
530 | func (h *Hub) LogInfof(format string, args ...interface{}) {
531 | 	h.logger.Infof(format, args...)
532 | }
533 | 
534 | // LogPrintf is like LogPrint but allows specifying a message format.
535 | func (h *Hub) LogPrintf(format string, args ...interface{}) {
536 | 	h.logger.Printf(format, args...)
537 | }
538 | 
539 | // LogWarnf is like LogWarn but allows specifying a message format.
540 | func (h *Hub) LogWarnf(format string, args ...interface{}) {
541 | 	h.logger.Warnf(format, args...)
542 | }
543 | 
544 | // LogErrorf is like LogError but allows specifying a message format.
545 | func (h *Hub) LogErrorf(format string, args ...interface{}) {
546 | 	h.logger.Errorf(format, args...)
547 | }
548 | 
549 | // LogFatalf is like LogFatal but allows specifying a message format.
550 | func (h *Hub) LogFatalf(format string, args ...interface{}) {
551 | 	h.logger.Fatalf(format, args...)
552 | }
553 | 
554 | // LogPanicf is like LogPanic but allows specifying a message format.
555 | func (h *Hub) LogPanicf(format string, args ...interface{}) {
556 | 	h.logger.Panicf(format, args...)
557 | }
558 | 
559 | func (h *Hub) connectClient(client WebsocketClient, response http.ResponseWriter, request *http.Request) error {
560 | 	h.clients.Set(client)
561 | 
562 | 	err := client.Connect(response, request)
563 | 	if err != nil {
564 | 		_ = h.clients.Unset(client.ID())
565 | 
566 | 		return errors.WithStack(err)
567 | 	}
568 | 
569 | 	connectHandler := h.connectHandler.Load().(ConnectHandler)
570 | 	connectHandler(client.ID())
571 | 
572 | 	return nil
573 | }
574 | 
575 | func (h *Hub) disconnectClient(client WebsocketClient) error {
576 | 	err := h.clients.Unset(client.ID())
577 | 	if err != nil {
578 | 		return errors.WithStack(err)
579 | 	}
580 | 
581 | 	err = client.Close()
582 | 	if err != nil {
583 | 		return errors.WithStack(err)
584 | 	}
585 | 
586 | 	disconnectHandler := h.disconnectHandler.Load().(DisconnectHandler)
587 | 	disconnectHandler(client.ID())
588 | 
589 | 	return nil
590 | }
591 | 
592 | func (h *Hub) wrapErrorHandler(handler ErrorHandler) ErrorHandler {
593 | 	return func(clientID UUID, err error) {
594 | 		handler(clientID, err)
595 | 
596 | 		// We should disconnect the client
597 | 		// if it reported (called the error_handler) that
598 | 		// an error has occurred while reading or writing a websocket
599 | 		_ = h.Disconnect(clientID)
600 | 	}
601 | }
602 | 
603 | // NewHub initializes a new Hub.
604 | func NewHub(
605 | 	options HubOptions,
606 | 	clientStore WebsocketClientStore,
607 | 	clientFactory WebsocketClientFactory,
608 | 	logger Logger,
609 | ) *Hub {
610 | 	hub := &Hub{
611 | 		options:       options,
612 | 		clients:       clientStore,
613 | 		clientFactory: clientFactory,
614 | 		logger:        logger,
615 | 		httpServer:    &http.Server{},
616 | 		httpServerTLS: &http.Server{},
617 | 	}
618 | 
619 | 	hub.connectHandler.Store(defaultConnectHandler)
620 | 	hub.disconnectHandler.Store(defaultDisconnectHandler)
621 | 	hub.receiveHandler.Store(defaultReceiveHandler)
622 | 	hub.errorHandler.Store(hub.wrapErrorHandler(defaultErrorHandler))
623 | 
624 | 	return hub
625 | }
626 | 
627 | // NewDefaultHub uses default dependencies to initializes a new hub.
628 | func NewDefaultHub() *Hub {
629 | 	logger := NewLogrusLogger(NewLogrusLoggerOptions())
630 | 
631 | 	upgraderOptions := NewGorillaConnectionUpgraderOptions()
632 | 	upgrader := NewGorillaConnectionUpgrader(upgraderOptions, logger)
633 | 
634 | 	clientStoreOptions := NewClientStoreOptions()
635 | 	clientStore := NewClientStore(clientStoreOptions, logger)
636 | 
637 | 	uuidGenerator := SatoriUUIDGenerator{}
638 | 
639 | 	clientOptions := NewClientOptions()
640 | 	clientFactory := NewClientFactory(clientOptions, uuidGenerator, upgrader, logger)
641 | 
642 | 	hubOptions := NewHubOptions()
643 | 	hub := NewHub(hubOptions, clientStore, clientFactory, logger)
644 | 
645 | 	return hub
646 | }
647 | 


--------------------------------------------------------------------------------
/hub_options.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import (
 4 | 	"time"
 5 | )
 6 | 
 7 | // HubOptions represents configuration of the hub.
 8 | type HubOptions struct {
 9 | 	// Time to gracefully shutdown a server
10 | 	ShutdownTimeout time.Duration
11 | 
12 | 	// Enable/disable debug mode.
13 | 	IsDebug bool
14 | 
15 | 	// Function execution time limit in debug mode.
16 | 	// Exceeding this time limit will cause a new warn log message.
17 | 	DebugFuncTimeLimit time.Duration
18 | }
19 | 
20 | // NewHubOptions initializes a new HubOptions.
21 | // nolint: gomnd
22 | func NewHubOptions() HubOptions {
23 | 	return HubOptions{
24 | 		ShutdownTimeout:    10 * time.Second,
25 | 		IsDebug:            false,
26 | 		DebugFuncTimeLimit: 1 * time.Millisecond,
27 | 	}
28 | }
29 | 


--------------------------------------------------------------------------------
/hub_options_test.go:
--------------------------------------------------------------------------------
 1 | package wspubsub_test
 2 | 
 3 | import (
 4 | 	"testing"
 5 | 
 6 | 	"github.com/kpeu3i/wspubsub"
 7 | 	"github.com/stretchr/testify/require"
 8 | )
 9 | 
10 | func TestNewHubOptions(t *testing.T) {
11 | 	options := wspubsub.NewHubOptions()
12 | 	require.NotZero(t, options.ShutdownTimeout)
13 | 	require.False(t, options.IsDebug)
14 | 	require.NotZero(t, options.DebugFuncTimeLimit)
15 | }
16 | 


--------------------------------------------------------------------------------
/hub_subscription_channel_required_error.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import (
 4 | 	"fmt"
 5 | 
 6 | 	"github.com/pkg/errors"
 7 | )
 8 | 
 9 | // HubSubscriptionChannelRequired returned when trying to subscribe with an empty channels list.
10 | type HubSubscriptionChannelRequired struct {
11 | 	message string
12 | }
13 | 
14 | // HubSubscriptionChannelRequired implements an error interface.
15 | func (e *HubSubscriptionChannelRequired) Error() string {
16 | 	return fmt.Sprintf("wspubsub: %s", e.message)
17 | }
18 | 
19 | // NewHubSubscriptionChannelRequiredError initializes a new HubSubscriptionChannelRequired.
20 | func NewHubSubscriptionChannelRequiredError() *HubSubscriptionChannelRequired {
21 | 	return &HubSubscriptionChannelRequired{message: "at least one subscription channel is required"}
22 | }
23 | 
24 | // IsHubSubscriptionChannelRequiredError checks if error type is HubSubscriptionChannelRequired.
25 | func IsHubSubscriptionChannelRequiredError(err error) (*HubSubscriptionChannelRequired, bool) {
26 | 	v, ok := errors.Cause(err).(*HubSubscriptionChannelRequired)
27 | 
28 | 	return v, ok
29 | }
30 | 


--------------------------------------------------------------------------------
/hub_subscription_channel_required_error_test.go:
--------------------------------------------------------------------------------
 1 | package wspubsub_test
 2 | 
 3 | import (
 4 | 	"errors"
 5 | 	"testing"
 6 | 
 7 | 	"github.com/kpeu3i/wspubsub"
 8 | 	"github.com/stretchr/testify/require"
 9 | )
10 | 
11 | func TestHubSubscriptionChannelRequiredError(t *testing.T) {
12 | 	rawErr := errors.New("TEST")
13 | 	err := wspubsub.NewHubSubscriptionChannelRequiredError()
14 | 	require.NotEmpty(t, clientID, err.Error())
15 | 
16 | 	e, ok := wspubsub.IsHubSubscriptionChannelRequiredError(err)
17 | 	require.NotNil(t, e)
18 | 	require.True(t, ok)
19 | 
20 | 	e, ok = wspubsub.IsHubSubscriptionChannelRequiredError(rawErr)
21 | 	require.Nil(t, e)
22 | 	require.False(t, ok)
23 | }
24 | 


--------------------------------------------------------------------------------
/hub_test.go:
--------------------------------------------------------------------------------
  1 | package wspubsub_test
  2 | 
  3 | import (
  4 | 	"net/http"
  5 | 	"net/http/httptest"
  6 | 	"sort"
  7 | 	"strings"
  8 | 	"testing"
  9 | 
 10 | 	"github.com/golang/mock/gomock"
 11 | 	"github.com/kpeu3i/wspubsub"
 12 | 	"github.com/kpeu3i/wspubsub/mock"
 13 | 	"github.com/pkg/errors"
 14 | 	"github.com/stretchr/testify/require"
 15 | )
 16 | 
 17 | func TestHub_Subscription(t *testing.T) {
 18 | 	ctrl := gomock.NewController(t)
 19 | 	defer ctrl.Finish()
 20 | 
 21 | 	logger := mock.NewMockLogger(ctrl)
 22 | 	clientStore := mock.NewMockWebsocketClientStore(ctrl)
 23 | 	clientFactory := mock.NewMockWebsocketClientFactory(ctrl)
 24 | 
 25 | 	channels := []string{"X", "Y", "W", "Z"}
 26 | 	unsubscribedChannels := []string{"X", "Y"}
 27 | 
 28 | 	clientStore.
 29 | 		EXPECT().
 30 | 		CountChannels(gomock.Any()).
 31 | 		Times(2).
 32 | 		DoAndReturn(func(cid wspubsub.UUID) (int, error) {
 33 | 			if cid == clientID {
 34 | 				return 1, nil
 35 | 			}
 36 | 
 37 | 			return 0, errors.New("count_channels_error")
 38 | 		})
 39 | 
 40 | 	clientStore.
 41 | 		EXPECT().
 42 | 		SetChannels(gomock.Eq(clientID), gomock.Eq(channels)).
 43 | 		Times(1)
 44 | 
 45 | 	clientStore.
 46 | 		EXPECT().
 47 | 		SetChannels(gomock.Eq(clientID), gomock.Eq(channels)).
 48 | 		Times(1).
 49 | 		Return(errors.New("set_channels_error"))
 50 | 
 51 | 	clientStore.
 52 | 		EXPECT().
 53 | 		UnsetChannels(gomock.Eq(clientID)).
 54 | 		Times(1)
 55 | 
 56 | 	clientStore.
 57 | 		EXPECT().
 58 | 		UnsetChannels(gomock.Eq(clientID), gomock.Eq(unsubscribedChannels)).
 59 | 		Times(1)
 60 | 
 61 | 	clientStore.
 62 | 		EXPECT().
 63 | 		UnsetChannels(gomock.Eq(clientID), gomock.Eq(unsubscribedChannels)).
 64 | 		Times(1).
 65 | 		Return(errors.New("unset_channels_error"))
 66 | 
 67 | 	hubOptions := wspubsub.NewHubOptions()
 68 | 	hub := wspubsub.NewHub(hubOptions, clientStore, clientFactory, logger)
 69 | 
 70 | 	t.Run("Subscribe to empty channels", func(t *testing.T) {
 71 | 		err := hub.Subscribe(clientID)
 72 | 		require.Equal(t, wspubsub.NewHubSubscriptionChannelRequiredError(), errors.Cause(err).(*wspubsub.HubSubscriptionChannelRequired))
 73 | 	})
 74 | 
 75 | 	t.Run("Subscribe to non-empty channels", func(t *testing.T) {
 76 | 		err := hub.Subscribe(clientID, channels...)
 77 | 		require.NoError(t, err)
 78 | 
 79 | 		err = hub.Subscribe(clientID, channels...)
 80 | 		require.Error(t, err)
 81 | 	})
 82 | 
 83 | 	t.Run("Unsubscribe from empty channels", func(t *testing.T) {
 84 | 		err := hub.Unsubscribe(clientID)
 85 | 		require.NoError(t, err)
 86 | 	})
 87 | 
 88 | 	t.Run("Unsubscribe from channels", func(t *testing.T) {
 89 | 		err := hub.Unsubscribe(clientID, unsubscribedChannels...)
 90 | 		require.NoError(t, err)
 91 | 
 92 | 		err = hub.Unsubscribe(clientID, unsubscribedChannels...)
 93 | 		require.Error(t, err)
 94 | 	})
 95 | 
 96 | 	t.Run("Check is subscribed", func(t *testing.T) {
 97 | 		isSubscribed := hub.IsSubscribed(clientID)
 98 | 		require.True(t, isSubscribed)
 99 | 
100 | 		isSubscribed = hub.IsSubscribed(wspubsub.UUID{})
101 | 		require.False(t, isSubscribed)
102 | 	})
103 | }
104 | 
105 | func TestHub_Channels(t *testing.T) {
106 | 	ctrl := gomock.NewController(t)
107 | 	defer ctrl.Finish()
108 | 
109 | 	logger := mock.NewMockLogger(ctrl)
110 | 	clientStore := mock.NewMockWebsocketClientStore(ctrl)
111 | 	clientFactory := mock.NewMockWebsocketClientFactory(ctrl)
112 | 
113 | 	channels := []string{"X", "Y", "W", "Z"}
114 | 	sort.Strings(channels)
115 | 
116 | 	clientStore.
117 | 		EXPECT().
118 | 		Channels(gomock.Any()).
119 | 		Times(2).
120 | 		DoAndReturn(func(cid wspubsub.UUID) ([]string, error) {
121 | 			if cid == clientID {
122 | 				return channels, nil
123 | 			}
124 | 
125 | 			return nil, errors.New("channels_error")
126 | 		})
127 | 
128 | 	hubOptions := wspubsub.NewHubOptions()
129 | 	hub := wspubsub.NewHub(hubOptions, clientStore, clientFactory, logger)
130 | 
131 | 	t.Run("Getting channels success", func(t *testing.T) {
132 | 		clientChannels, err := hub.Channels(clientID)
133 | 		sort.Strings(clientChannels)
134 | 		require.NoError(t, err)
135 | 		require.Equal(t, channels, clientChannels)
136 | 	})
137 | 
138 | 	t.Run("Getting channels error", func(t *testing.T) {
139 | 		clientChannels, err := hub.Channels(wspubsub.UUID{})
140 | 		sort.Strings(clientChannels)
141 | 		require.Error(t, err)
142 | 		require.Nil(t, clientChannels)
143 | 	})
144 | }
145 | 
146 | func TestHub_Count(t *testing.T) {
147 | 	ctrl := gomock.NewController(t)
148 | 	defer ctrl.Finish()
149 | 
150 | 	logger := mock.NewMockLogger(ctrl)
151 | 	clientStore := mock.NewMockWebsocketClientStore(ctrl)
152 | 	clientFactory := mock.NewMockWebsocketClientFactory(ctrl)
153 | 
154 | 	count := 100
155 | 	channels := []string{"X", "Y", "W", "Z"}
156 | 	sort.Strings(channels)
157 | 
158 | 	clientStore.
159 | 		EXPECT().
160 | 		Count(gomock.Eq(channels)).
161 | 		Times(1).
162 | 		Return(count)
163 | 
164 | 	hubOptions := wspubsub.NewHubOptions()
165 | 	hub := wspubsub.NewHub(hubOptions, clientStore, clientFactory, logger)
166 | 
167 | 	require.Equal(t, count, hub.Count(channels...))
168 | }
169 | 
170 | func TestHub_Publish(t *testing.T) {
171 | 	ctrl := gomock.NewController(t)
172 | 	defer ctrl.Finish()
173 | 
174 | 	logger := mock.NewMockLogger(ctrl)
175 | 	clientStore := mock.NewMockWebsocketClientStore(ctrl)
176 | 	clientFactory := mock.NewMockWebsocketClientFactory(ctrl)
177 | 	client := mock.NewMockWebsocketClient(ctrl)
178 | 
179 | 	message := wspubsub.NewTextMessageFromString("TEST")
180 | 	findErrText := "find_error"
181 | 	sendErrText := "send_error"
182 | 
183 | 	clientStore.
184 | 		EXPECT().
185 | 		Find(gomock.Any(), gomock.Any()).
186 | 		Times(2).
187 | 		DoAndReturn(func(fn wspubsub.IterateFunc, channels ...string) error {
188 | 			return fn(client)
189 | 		})
190 | 
191 | 	clientStore.
192 | 		EXPECT().
193 | 		Find(gomock.Any(), gomock.Any()).
194 | 		Times(1).
195 | 		DoAndReturn(func(fn wspubsub.IterateFunc, channels ...string) error {
196 | 			return errors.New(findErrText)
197 | 		})
198 | 
199 | 	clientStore.
200 | 		EXPECT().
201 | 		Unset(gomock.Eq(clientID)).
202 | 		Times(1)
203 | 
204 | 	client.
205 | 		EXPECT().
206 | 		Send(gomock.Eq(message)).
207 | 		Times(1)
208 | 
209 | 	client.
210 | 		EXPECT().
211 | 		Send(gomock.Eq(message)).
212 | 		Times(1).
213 | 		Return(errors.New(sendErrText))
214 | 
215 | 	client.
216 | 		EXPECT().
217 | 		ID().
218 | 		AnyTimes().
219 | 		Return(clientID)
220 | 
221 | 	client.
222 | 		EXPECT().
223 | 		Close().
224 | 		Times(1)
225 | 
226 | 	hubOptions := wspubsub.NewHubOptions()
227 | 	hub := wspubsub.NewHub(hubOptions, clientStore, clientFactory, logger)
228 | 
229 | 	t.Run("Publishing message success", func(t *testing.T) {
230 | 		numClients, err := hub.Publish(message)
231 | 		require.NoError(t, err)
232 | 		require.Equal(t, 1, numClients)
233 | 
234 | 		numClients, err = hub.Publish(message)
235 | 		require.NoError(t, err)
236 | 		require.Equal(t, 0, numClients)
237 | 	})
238 | 
239 | 	t.Run("Publishing message error", func(t *testing.T) {
240 | 		numClients, err := hub.Publish(message)
241 | 		require.Error(t, err)
242 | 		require.True(t, strings.Contains(err.Error(), findErrText))
243 | 		require.Equal(t, 0, numClients)
244 | 	})
245 | }
246 | 
247 | func TestHub_Send(t *testing.T) {
248 | 	ctrl := gomock.NewController(t)
249 | 	defer ctrl.Finish()
250 | 
251 | 	logger := mock.NewMockLogger(ctrl)
252 | 	clientStore := mock.NewMockWebsocketClientStore(ctrl)
253 | 	clientFactory := mock.NewMockWebsocketClientFactory(ctrl)
254 | 	client := mock.NewMockWebsocketClient(ctrl)
255 | 
256 | 	message := wspubsub.NewTextMessageFromString("TEST")
257 | 	unsetErrText := "unset_error"
258 | 	closeErrText := "close_error"
259 | 
260 | 	// *** Success case
261 | 	clientStore.
262 | 		EXPECT().
263 | 		Get(gomock.Eq(clientID)).
264 | 		Times(1).
265 | 		Return(client, nil)
266 | 
267 | 	client.
268 | 		EXPECT().
269 | 		Send(gomock.Eq(message)).
270 | 		Times(1)
271 | 
272 | 	// *** Error case #1
273 | 	clientStore.
274 | 		EXPECT().
275 | 		Get(gomock.Eq(clientID)).
276 | 		Times(1).
277 | 		Return(client, nil)
278 | 
279 | 	client.
280 | 		EXPECT().
281 | 		Send(gomock.Eq(message)).
282 | 		Times(1).
283 | 		Return(errors.New(unsetErrText))
284 | 
285 | 	// *** Error case #2
286 | 	clientStore.
287 | 		EXPECT().
288 | 		Get(gomock.Eq(clientID)).
289 | 		Times(1).
290 | 		Return(nil, errors.New(closeErrText))
291 | 
292 | 	hubOptions := wspubsub.NewHubOptions()
293 | 	hub := wspubsub.NewHub(hubOptions, clientStore, clientFactory, logger)
294 | 
295 | 	t.Run("Sending message success", func(t *testing.T) {
296 | 		err := hub.Send(clientID, message)
297 | 		require.NoError(t, err)
298 | 	})
299 | 
300 | 	t.Run("Sending message error #1", func(t *testing.T) {
301 | 		err := hub.Send(clientID, message)
302 | 		require.Error(t, err)
303 | 		require.True(t, strings.Contains(err.Error(), unsetErrText))
304 | 	})
305 | 
306 | 	t.Run("Sending message error #2", func(t *testing.T) {
307 | 		err := hub.Send(clientID, message)
308 | 		require.Error(t, err)
309 | 		require.True(t, strings.Contains(err.Error(), closeErrText))
310 | 	})
311 | }
312 | 
313 | func TestHub_Disconnect(t *testing.T) {
314 | 	ctrl := gomock.NewController(t)
315 | 	defer ctrl.Finish()
316 | 
317 | 	logger := mock.NewMockLogger(ctrl)
318 | 	clientStore := mock.NewMockWebsocketClientStore(ctrl)
319 | 	clientFactory := mock.NewMockWebsocketClientFactory(ctrl)
320 | 	client := mock.NewMockWebsocketClient(ctrl)
321 | 
322 | 	unsetErrText := "unset_error"
323 | 	closeErrText := "close_error"
324 | 	getErrText := "get_error"
325 | 
326 | 	// *** Success case
327 | 	clientStore.
328 | 		EXPECT().
329 | 		Get(gomock.Eq(clientID)).
330 | 		Times(1).
331 | 		Return(client, nil)
332 | 
333 | 	clientStore.
334 | 		EXPECT().
335 | 		Unset(gomock.Eq(clientID)).
336 | 		Times(1)
337 | 
338 | 	client.
339 | 		EXPECT().
340 | 		ID().
341 | 		AnyTimes().
342 | 		Return(clientID)
343 | 
344 | 	client.
345 | 		EXPECT().
346 | 		Close().
347 | 		Times(1)
348 | 
349 | 	// *** Error case #1
350 | 	clientStore.
351 | 		EXPECT().
352 | 		Get(gomock.Eq(clientID)).
353 | 		Times(1).
354 | 		Return(client, nil)
355 | 
356 | 	clientStore.
357 | 		EXPECT().
358 | 		Unset(gomock.Eq(clientID)).
359 | 		Times(1).
360 | 		Return(errors.New(unsetErrText))
361 | 
362 | 	// *** Error case #2
363 | 	clientStore.
364 | 		EXPECT().
365 | 		Get(gomock.Eq(clientID)).
366 | 		Times(1).
367 | 		Return(client, nil)
368 | 
369 | 	clientStore.
370 | 		EXPECT().
371 | 		Unset(gomock.Eq(clientID)).
372 | 		Times(1)
373 | 
374 | 	client.
375 | 		EXPECT().
376 | 		Close().
377 | 		Times(1).
378 | 		Return(errors.New(closeErrText))
379 | 
380 | 	// *** Error case #3
381 | 	clientStore.
382 | 		EXPECT().
383 | 		Get(gomock.Eq(clientID)).
384 | 		Times(1).
385 | 		Return(nil, errors.New(getErrText))
386 | 
387 | 	hubOptions := wspubsub.NewHubOptions()
388 | 	hub := wspubsub.NewHub(hubOptions, clientStore, clientFactory, logger)
389 | 
390 | 	t.Run("Disconnection success", func(t *testing.T) {
391 | 		err := hub.Disconnect(clientID)
392 | 		require.NoError(t, err)
393 | 	})
394 | 
395 | 	t.Run("Disconnection error #1", func(t *testing.T) {
396 | 		err := hub.Disconnect(clientID)
397 | 		require.Error(t, err)
398 | 		require.True(t, strings.Contains(err.Error(), unsetErrText))
399 | 	})
400 | 
401 | 	t.Run("Disconnection error #2", func(t *testing.T) {
402 | 		err := hub.Disconnect(clientID)
403 | 		require.Error(t, err)
404 | 		require.True(t, strings.Contains(err.Error(), closeErrText))
405 | 	})
406 | 
407 | 	t.Run("Disconnection error #3", func(t *testing.T) {
408 | 		err := hub.Disconnect(clientID)
409 | 		require.Error(t, err)
410 | 		require.True(t, strings.Contains(err.Error(), getErrText))
411 | 	})
412 | }
413 | 
414 | func TestHub_ConnectSuccess(t *testing.T) {
415 | 	ctrl := gomock.NewController(t)
416 | 	defer ctrl.Finish()
417 | 
418 | 	logger := mock.NewMockLogger(ctrl)
419 | 
420 | 	logger.
421 | 		EXPECT().
422 | 		Infof(gomock.Any(), gomock.Any()).
423 | 		AnyTimes()
424 | 
425 | 	logger.
426 | 		EXPECT().
427 | 		Info(gomock.Any()).
428 | 		AnyTimes()
429 | 
430 | 	clientStore := mock.NewMockWebsocketClientStore(ctrl)
431 | 	clientFactory := mock.NewMockWebsocketClientFactory(ctrl)
432 | 	client := mock.NewMockWebsocketClient(ctrl)
433 | 
434 | 	connectHandlerNumCalls := 0
435 | 	connectHandler := func(cid wspubsub.UUID) {
436 | 		require.Equal(t, clientID, cid)
437 | 		connectHandlerNumCalls++
438 | 	}
439 | 
440 | 	disconnectHandler := func(cid wspubsub.UUID) {
441 | 		t.Error("Unexpected call of: disconnect_handler")
442 | 	}
443 | 
444 | 	receiveHandler := func(cid wspubsub.UUID, message wspubsub.Message) {
445 | 		t.Error("Unexpected call of: receive_handler")
446 | 	}
447 | 
448 | 	errorHandler := func(cid wspubsub.UUID, err error) {
449 | 		t.Error("Unexpected call of: error_handler")
450 | 	}
451 | 
452 | 	request := httptest.NewRequest("GET", "/", nil)
453 | 	response := httptest.NewRecorder()
454 | 
455 | 	hubOptions := wspubsub.NewHubOptions()
456 | 	hub := wspubsub.NewHub(hubOptions, clientStore, clientFactory, logger)
457 | 	hub.OnConnect(connectHandler)
458 | 	hub.OnDisconnect(disconnectHandler)
459 | 	hub.OnReceive(receiveHandler)
460 | 	hub.OnError(errorHandler)
461 | 
462 | 	clientFactory.
463 | 		EXPECT().
464 | 		Create().
465 | 		Times(1).
466 | 		Return(client)
467 | 
468 | 	clientStore.
469 | 		EXPECT().
470 | 		Set(gomock.Eq(client)).
471 | 		Times(1)
472 | 
473 | 	client.
474 | 		EXPECT().
475 | 		OnReceive(gomock.Any()).
476 | 		Times(1)
477 | 
478 | 	client.
479 | 		EXPECT().
480 | 		OnError(gomock.Any()).
481 | 		Times(1)
482 | 
483 | 	client.
484 | 		EXPECT().
485 | 		ID().
486 | 		AnyTimes().
487 | 		Return(clientID)
488 | 
489 | 	client.
490 | 		EXPECT().
491 | 		Connect(gomock.Eq(response), gomock.Eq(request)).
492 | 		Times(1)
493 | 
494 | 	hub.ServeHTTP(response, request)
495 | 
496 | 	require.Equal(t, http.StatusOK, response.Result().StatusCode)
497 | 	require.Equal(t, 1, connectHandlerNumCalls)
498 | }
499 | 
500 | func TestHub_ConnectError(t *testing.T) {
501 | 	ctrl := gomock.NewController(t)
502 | 	defer ctrl.Finish()
503 | 
504 | 	logger := mock.NewMockLogger(ctrl)
505 | 
506 | 	logger.
507 | 		EXPECT().
508 | 		Infof(gomock.Any(), gomock.Any()).
509 | 		AnyTimes()
510 | 
511 | 	logger.
512 | 		EXPECT().
513 | 		Info(gomock.Any()).
514 | 		AnyTimes()
515 | 
516 | 	clientStore := mock.NewMockWebsocketClientStore(ctrl)
517 | 	clientFactory := mock.NewMockWebsocketClientFactory(ctrl)
518 | 	client := mock.NewMockWebsocketClient(ctrl)
519 | 
520 | 	connectErrText := "connect_error"
521 | 
522 | 	connectHandler := func(cid wspubsub.UUID) {
523 | 		t.Error("Unexpected call of: connect_handler")
524 | 	}
525 | 
526 | 	disconnectHandler := func(cid wspubsub.UUID) {
527 | 		t.Error("Unexpected call of: disconnect_handler")
528 | 	}
529 | 
530 | 	receiveHandler := func(cid wspubsub.UUID, message wspubsub.Message) {
531 | 		t.Error("Unexpected call of: receive_handler")
532 | 	}
533 | 
534 | 	errorHandler := func(cid wspubsub.UUID, err error) {
535 | 		t.Error("Unexpected call of: error_handler")
536 | 	}
537 | 
538 | 	request := httptest.NewRequest("GET", "/", nil)
539 | 	response := httptest.NewRecorder()
540 | 
541 | 	hubOptions := wspubsub.NewHubOptions()
542 | 	hub := wspubsub.NewHub(hubOptions, clientStore, clientFactory, logger)
543 | 	hub.OnConnect(connectHandler)
544 | 	hub.OnDisconnect(disconnectHandler)
545 | 	hub.OnReceive(receiveHandler)
546 | 	hub.OnError(errorHandler)
547 | 
548 | 	clientFactory.
549 | 		EXPECT().
550 | 		Create().
551 | 		Times(1).
552 | 		Return(client)
553 | 
554 | 	clientStore.
555 | 		EXPECT().
556 | 		Set(gomock.Eq(client)).
557 | 		Times(1)
558 | 
559 | 	clientStore.
560 | 		EXPECT().
561 | 		Unset(gomock.Eq(clientID)).
562 | 		Times(1)
563 | 
564 | 	client.
565 | 		EXPECT().
566 | 		OnReceive(gomock.Any()).
567 | 		Times(1)
568 | 
569 | 	client.
570 | 		EXPECT().
571 | 		OnError(gomock.Any()).
572 | 		Times(1)
573 | 
574 | 	client.
575 | 		EXPECT().
576 | 		ID().
577 | 		AnyTimes().
578 | 		Return(clientID)
579 | 
580 | 	client.
581 | 		EXPECT().
582 | 		Connect(gomock.Eq(response), gomock.Eq(request)).
583 | 		Times(1).
584 | 		Return(errors.New(connectErrText))
585 | 
586 | 	hub.ServeHTTP(response, request)
587 | 
588 | 	require.Equal(t, http.StatusInternalServerError, response.Result().StatusCode)
589 | }
590 | 
591 | func TestHub_ErrorHandler(t *testing.T) {
592 | 	ctrl := gomock.NewController(t)
593 | 	defer ctrl.Finish()
594 | 
595 | 	logger := mock.NewMockLogger(ctrl)
596 | 
597 | 	logger.
598 | 		EXPECT().
599 | 		Infof(gomock.Any(), gomock.Any()).
600 | 		AnyTimes()
601 | 
602 | 	logger.
603 | 		EXPECT().
604 | 		Info(gomock.Any()).
605 | 		AnyTimes()
606 | 
607 | 	clientStore := mock.NewMockWebsocketClientStore(ctrl)
608 | 	clientFactory := mock.NewMockWebsocketClientFactory(ctrl)
609 | 	client := mock.NewMockWebsocketClient(ctrl)
610 | 
611 | 	request := httptest.NewRequest("GET", "/", nil)
612 | 	response := httptest.NewRecorder()
613 | 
614 | 	var errorHandler wspubsub.ErrorHandler
615 | 	errText := "some_error"
616 | 
617 | 	clientFactory.
618 | 		EXPECT().
619 | 		Create().
620 | 		Times(1).
621 | 		Return(client)
622 | 
623 | 	clientStore.
624 | 		EXPECT().
625 | 		Set(gomock.Eq(client)).
626 | 		Times(1)
627 | 
628 | 	clientStore.
629 | 		EXPECT().
630 | 		Unset(gomock.Eq(clientID)).
631 | 		Times(1)
632 | 
633 | 	clientStore.
634 | 		EXPECT().
635 | 		Get(gomock.Eq(clientID)).
636 | 		Times(1).
637 | 		Return(client, nil)
638 | 
639 | 	client.
640 | 		EXPECT().
641 | 		ID().
642 | 		AnyTimes().
643 | 		Return(clientID)
644 | 
645 | 	client.
646 | 		EXPECT().
647 | 		OnReceive(gomock.Any()).
648 | 		Times(1)
649 | 
650 | 	client.
651 | 		EXPECT().
652 | 		OnError(gomock.Any()).
653 | 		Times(1).
654 | 		DoAndReturn(func(handler func(cid wspubsub.UUID, err error)) {
655 | 			// Just remember the error handler to call it later
656 | 			errorHandler = handler
657 | 		})
658 | 
659 | 	client.
660 | 		EXPECT().
661 | 		Connect(gomock.Eq(response), gomock.Eq(request)).
662 | 		Times(1)
663 | 
664 | 	client.
665 | 		EXPECT().
666 | 		Close().
667 | 		Times(1)
668 | 
669 | 	hubOptions := wspubsub.NewHubOptions()
670 | 	hub := wspubsub.NewHub(hubOptions, clientStore, clientFactory, logger)
671 | 	hub.OnError(func(cid wspubsub.UUID, err error) {
672 | 		require.True(t, strings.Contains(err.Error(), errText))
673 | 	})
674 | 	hub.ServeHTTP(response, request)
675 | 
676 | 	errorHandler(clientID, errors.New(errText))
677 | 	require.Equal(t, http.StatusOK, response.Result().StatusCode)
678 | }
679 | 
680 | func TestHub_Close(t *testing.T) {
681 | 	ctrl := gomock.NewController(t)
682 | 	defer ctrl.Finish()
683 | 
684 | 	logger := mock.NewMockLogger(ctrl)
685 | 
686 | 	logger.
687 | 		EXPECT().
688 | 		Infof(gomock.Any(), gomock.Any()).
689 | 		AnyTimes()
690 | 
691 | 	logger.
692 | 		EXPECT().
693 | 		Info(gomock.Any()).
694 | 		AnyTimes()
695 | 
696 | 	clientStore := mock.NewMockWebsocketClientStore(ctrl)
697 | 	clientFactory := mock.NewMockWebsocketClientFactory(ctrl)
698 | 	client := mock.NewMockWebsocketClient(ctrl)
699 | 
700 | 	hubOptions := wspubsub.NewHubOptions()
701 | 	hub := wspubsub.NewHub(hubOptions, clientStore, clientFactory, logger)
702 | 
703 | 	findErrText := "find_error"
704 | 
705 | 	// *** Success case
706 | 	client.
707 | 		EXPECT().
708 | 		ID().
709 | 		AnyTimes().
710 | 		Return(clientID)
711 | 
712 | 	client.
713 | 		EXPECT().
714 | 		Close().
715 | 		Times(1)
716 | 
717 | 	clientStore.
718 | 		EXPECT().
719 | 		Unset(gomock.Eq(clientID)).
720 | 		Times(1)
721 | 
722 | 	clientStore.
723 | 		EXPECT().
724 | 		Find(gomock.Any(), gomock.Eq([]string{})).
725 | 		Times(1).
726 | 		DoAndReturn(func(fn wspubsub.IterateFunc, channels ...string) error {
727 | 			err := fn(client)
728 | 			require.NoError(t, err)
729 | 
730 | 			return nil
731 | 		})
732 | 
733 | 	// *** Error case
734 | 	clientStore.
735 | 		EXPECT().
736 | 		Find(gomock.Any(), gomock.Eq([]string{})).
737 | 		Times(1).
738 | 		DoAndReturn(func(fn wspubsub.IterateFunc, channels ...string) error {
739 | 			return errors.New(findErrText)
740 | 		})
741 | 
742 | 	t.Run("Closing success", func(t *testing.T) {
743 | 		err := hub.Close()
744 | 		require.NoError(t, err)
745 | 	})
746 | 
747 | 	t.Run("Closing error", func(t *testing.T) {
748 | 		err := hub.Close()
749 | 		require.Error(t, err)
750 | 		require.True(t, strings.Contains(err.Error(), findErrText))
751 | 	})
752 | }
753 | 
754 | func TestHub_Logging(t *testing.T) {
755 | 	ctrl := gomock.NewController(t)
756 | 	defer ctrl.Finish()
757 | 
758 | 	logger := mock.NewMockLogger(ctrl)
759 | 
760 | 	format := "log_format"
761 | 	message := "log_message"
762 | 
763 | 	logger.EXPECT().Debug(gomock.Eq(message)).Times(1)
764 | 	logger.EXPECT().Info(gomock.Eq(message)).Times(1)
765 | 	logger.EXPECT().Print(gomock.Eq(message)).Times(1)
766 | 	logger.EXPECT().Warn(gomock.Eq(message)).Times(1)
767 | 	logger.EXPECT().Error(gomock.Eq(message)).Times(1)
768 | 	logger.EXPECT().Fatal(gomock.Eq(message)).Times(1)
769 | 	logger.EXPECT().Panic(gomock.Eq(message)).Times(1)
770 | 
771 | 	logger.EXPECT().Debugln(gomock.Eq(message)).Times(1)
772 | 	logger.EXPECT().Infoln(gomock.Eq(message)).Times(1)
773 | 	logger.EXPECT().Println(gomock.Eq(message)).Times(1)
774 | 	logger.EXPECT().Warnln(gomock.Eq(message)).Times(1)
775 | 	logger.EXPECT().Errorln(gomock.Eq(message)).Times(1)
776 | 	logger.EXPECT().Fatalln(gomock.Eq(message)).Times(1)
777 | 	logger.EXPECT().Panicln(gomock.Eq(message)).Times(1)
778 | 
779 | 	logger.EXPECT().Debugf(gomock.Eq(format), gomock.Eq(message)).Times(1)
780 | 	logger.EXPECT().Infof(gomock.Eq(format), gomock.Eq(message)).Times(1)
781 | 	logger.EXPECT().Printf(gomock.Eq(format), gomock.Eq(message)).Times(1)
782 | 	logger.EXPECT().Warnf(gomock.Eq(format), gomock.Eq(message)).Times(1)
783 | 	logger.EXPECT().Errorf(gomock.Eq(format), gomock.Eq(message)).Times(1)
784 | 	logger.EXPECT().Fatalf(gomock.Eq(format), gomock.Eq(message)).Times(1)
785 | 	logger.EXPECT().Panicf(gomock.Eq(format), gomock.Eq(message)).Times(1)
786 | 
787 | 	clientStore := mock.NewMockWebsocketClientStore(ctrl)
788 | 	clientFactory := mock.NewMockWebsocketClientFactory(ctrl)
789 | 
790 | 	hubOptions := wspubsub.NewHubOptions()
791 | 	hub := wspubsub.NewHub(hubOptions, clientStore, clientFactory, logger)
792 | 
793 | 	hub.LogDebug(message)
794 | 	hub.LogInfo(message)
795 | 	hub.LogPrint(message)
796 | 	hub.LogWarn(message)
797 | 	hub.LogError(message)
798 | 	hub.LogFatal(message)
799 | 	hub.LogPanic(message)
800 | 
801 | 	hub.LogDebugln(message)
802 | 	hub.LogInfoln(message)
803 | 	hub.LogPrintln(message)
804 | 	hub.LogWarnln(message)
805 | 	hub.LogErrorln(message)
806 | 	hub.LogFatalln(message)
807 | 	hub.LogPanicln(message)
808 | 
809 | 	hub.LogDebugf(format, message)
810 | 	hub.LogInfof(format, message)
811 | 	hub.LogPrintf(format, message)
812 | 	hub.LogWarnf(format, message)
813 | 	hub.LogErrorf(format, message)
814 | 	hub.LogFatalf(format, message)
815 | 	hub.LogPanicf(format, message)
816 | }
817 | 


--------------------------------------------------------------------------------
/logrus_logger.go:
--------------------------------------------------------------------------------
  1 | package wspubsub
  2 | 
  3 | import (
  4 | 	"github.com/sirupsen/logrus"
  5 | )
  6 | 
  7 | // LogrusFormatter enumerates possible formatters.
  8 | type LogrusFormatter uint32
  9 | 
 10 | const (
 11 | 	// LogrusFormatterText formats logs into text.
 12 | 	LogrusFormatterText LogrusFormatter = iota
 13 | 
 14 | 	// LogrusFormatterJSON formats logs into parsable JSON.
 15 | 	LogrusFormatterJSON
 16 | )
 17 | 
 18 | // LogrusLevel enumerates possible logging levels.
 19 | type LogrusLevel uint32
 20 | 
 21 | const (
 22 | 	LogrusLevelPanic LogrusLevel = iota
 23 | 	LogrusLevelFatal
 24 | 	LogrusLevelError
 25 | 	LogrusLevelWarn
 26 | 	LogrusLevelInfo
 27 | 	LogrusLevelDebug
 28 | 	LogrusLevelTrace
 29 | )
 30 | 
 31 | var _ Logger = (*LogrusLogger)(nil)
 32 | 
 33 | // LogrusLogger is an implementation of Logger.
 34 | type LogrusLogger struct {
 35 | 	logger *logrus.Logger
 36 | }
 37 | 
 38 | // Debug logs a message at level Debug.
 39 | func (l LogrusLogger) Debug(args ...interface{}) {
 40 | 	l.logger.Debug(args...)
 41 | }
 42 | 
 43 | // Info logs a message at level Info.
 44 | func (l LogrusLogger) Info(args ...interface{}) {
 45 | 	l.logger.Info(args...)
 46 | }
 47 | 
 48 | // Print logs a message at level Info.
 49 | func (l LogrusLogger) Print(args ...interface{}) {
 50 | 	l.logger.Print(args...)
 51 | }
 52 | 
 53 | // Warn logs a message at level Warn.
 54 | func (l LogrusLogger) Warn(args ...interface{}) {
 55 | 	l.logger.Warn(args...)
 56 | }
 57 | 
 58 | // Error logs a message at level Error.
 59 | func (l LogrusLogger) Error(args ...interface{}) {
 60 | 	l.logger.Error(args...)
 61 | }
 62 | 
 63 | // Fatal logs a message at level Fatal then the process will exit with status set to 1.
 64 | func (l LogrusLogger) Fatal(args ...interface{}) {
 65 | 	l.logger.Fatal(args...)
 66 | }
 67 | 
 68 | // Panic logs a message at level Panic and panics.
 69 | func (l LogrusLogger) Panic(args ...interface{}) {
 70 | 	l.logger.Panic(args...)
 71 | }
 72 | 
 73 | // Debugln is like LogDebug but adds a new line.
 74 | func (l LogrusLogger) Debugln(args ...interface{}) {
 75 | 	l.logger.Debugln(args...)
 76 | }
 77 | 
 78 | // Infoln is like LogInfo but adds a new line.
 79 | func (l LogrusLogger) Infoln(args ...interface{}) {
 80 | 	l.logger.Infoln(args...)
 81 | }
 82 | 
 83 | // Println is like LogPrint but adds a new line.
 84 | func (l LogrusLogger) Println(args ...interface{}) {
 85 | 	l.logger.Println(args...)
 86 | }
 87 | 
 88 | // Warnln is like LogWarn but adds a new line.
 89 | func (l LogrusLogger) Warnln(args ...interface{}) {
 90 | 	l.logger.Warnln(args...)
 91 | }
 92 | 
 93 | // Errorln is like LogError but adds a new line.
 94 | func (l LogrusLogger) Errorln(args ...interface{}) {
 95 | 	l.logger.Errorln(args...)
 96 | }
 97 | 
 98 | // Fatalln is like LogFatal but adds a new line.
 99 | func (l LogrusLogger) Fatalln(args ...interface{}) {
100 | 	l.logger.Fatalln(args...)
101 | }
102 | 
103 | // Panicln is like LogPanic but adds a new line.
104 | func (l LogrusLogger) Panicln(args ...interface{}) {
105 | 	l.logger.Panicln(args...)
106 | }
107 | 
108 | // Debugf is like LogDebug but allows specifying a message format.
109 | func (l LogrusLogger) Debugf(format string, args ...interface{}) {
110 | 	l.logger.Debugf(format, args...)
111 | }
112 | 
113 | // Infof is like LogInfo but allows specifying a message format.
114 | func (l LogrusLogger) Infof(format string, args ...interface{}) {
115 | 	l.logger.Infof(format, args...)
116 | }
117 | 
118 | // Printf is like LogPrint but allows specifying a message format.
119 | func (l LogrusLogger) Printf(format string, args ...interface{}) {
120 | 	l.logger.Printf(format, args...)
121 | }
122 | 
123 | // Warnf is like LogWarn but allows specifying a message format.
124 | func (l LogrusLogger) Warnf(format string, args ...interface{}) {
125 | 	l.logger.Warnf(format, args...)
126 | }
127 | 
128 | // Errorf is like LogError but allows specifying a message format.
129 | func (l LogrusLogger) Errorf(format string, args ...interface{}) {
130 | 	l.logger.Errorf(format, args...)
131 | }
132 | 
133 | // Fatalf is like LogFatal but allows specifying a message format.
134 | func (l LogrusLogger) Fatalf(format string, args ...interface{}) {
135 | 	l.logger.Fatalf(format, args...)
136 | }
137 | 
138 | // Panicf is like LogPanic but allows specifying a message format.
139 | func (l LogrusLogger) Panicf(format string, args ...interface{}) {
140 | 	l.logger.Panicf(format, args...)
141 | }
142 | 
143 | // NewLogrusLogger initializes a new LogrusLogger.
144 | func NewLogrusLogger(options LogrusLoggerOptions) *LogrusLogger {
145 | 	logrusLogger := logrus.New()
146 | 	logrusLogger.SetLevel(logrus.Level(options.Level))
147 | 	logrusLogger.SetOutput(options.Output)
148 | 
149 | 	switch options.Formatter {
150 | 	case LogrusFormatterText:
151 | 		logrusLogger.SetFormatter(&logrus.TextFormatter{})
152 | 	case LogrusFormatterJSON:
153 | 		logrusLogger.SetFormatter(&logrus.JSONFormatter{})
154 | 	}
155 | 
156 | 	return &LogrusLogger{logger: logrusLogger}
157 | }
158 | 


--------------------------------------------------------------------------------
/logrus_options.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import (
 4 | 	"io"
 5 | 	"os"
 6 | )
 7 | 
 8 | // LogrusLoggerOptions represents configuration of the LogrusLogger.
 9 | type LogrusLoggerOptions struct {
10 | 	Level     LogrusLevel
11 | 	Formatter LogrusFormatter
12 | 	Output    io.Writer
13 | }
14 | 
15 | // NewLogrusLoggerOptions initializes a new LogrusLoggerOptions.
16 | func NewLogrusLoggerOptions() LogrusLoggerOptions {
17 | 	options := LogrusLoggerOptions{
18 | 		Level:     LogrusLevelInfo,
19 | 		Formatter: LogrusFormatterText,
20 | 		Output:    os.Stdout,
21 | 	}
22 | 
23 | 	return options
24 | }
25 | 


--------------------------------------------------------------------------------
/logrus_options_test.go:
--------------------------------------------------------------------------------
 1 | package wspubsub_test
 2 | 
 3 | import (
 4 | 	"testing"
 5 | 
 6 | 	"github.com/kpeu3i/wspubsub"
 7 | 	"github.com/stretchr/testify/require"
 8 | )
 9 | 
10 | func TestNewLogrusOptions(t *testing.T) {
11 | 	options := wspubsub.NewLogrusLoggerOptions()
12 | 	require.NotZero(t, options.Output)
13 | }
14 | 


--------------------------------------------------------------------------------
/message.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | // MessageType enumerates possible message types.
 4 | type MessageType byte
 5 | 
 6 | const (
 7 | 	MessageTypeText   MessageType = 1
 8 | 	MessageTypeBinary MessageType = 2
 9 | 	MessageTypePing   MessageType = 9
10 | )
11 | 
12 | // Message represents a data type to send over a WebSocket connection.
13 | type Message struct {
14 | 	Type    MessageType
15 | 	Payload []byte
16 | }
17 | 
18 | // NewTextMessage initializes a new text Message from bytes.
19 | func NewTextMessage(payload []byte) Message {
20 | 	return Message{Type: MessageTypeText, Payload: payload}
21 | }
22 | 
23 | // NewTextMessage initializes a new text Message from string.
24 | func NewTextMessageFromString(payload string) Message {
25 | 	return NewTextMessage([]byte(payload))
26 | }
27 | 
28 | // NewBinaryMessage initializes a new binary Message from bytes.
29 | func NewBinaryMessage(payload []byte) Message {
30 | 	return Message{Type: MessageTypeBinary, Payload: payload}
31 | }
32 | 
33 | // NewBinaryMessageFromString initializes a new binary Message from string.
34 | func NewBinaryMessageFromString(payload string) Message {
35 | 	return NewBinaryMessage([]byte(payload))
36 | }
37 | 
38 | // NewPingMessage initializes a new ping Message.
39 | func NewPingMessage() Message {
40 | 	return Message{Type: MessageTypePing}
41 | }
42 | 


--------------------------------------------------------------------------------
/message_test.go:
--------------------------------------------------------------------------------
 1 | package wspubsub_test
 2 | 
 3 | import (
 4 | 	"testing"
 5 | 
 6 | 	"github.com/kpeu3i/wspubsub"
 7 | 	"github.com/stretchr/testify/require"
 8 | )
 9 | 
10 | func TestNewTextMessage(t *testing.T) {
11 | 	b := []byte("TEST")
12 | 	message := wspubsub.NewTextMessage(b)
13 | 	require.Equal(t, wspubsub.MessageTypeText, message.Type)
14 | 	require.Equal(t, b, message.Payload)
15 | 
16 | 	s := "TEST"
17 | 	message = wspubsub.NewTextMessageFromString(s)
18 | 	require.Equal(t, wspubsub.MessageTypeText, message.Type)
19 | 	require.Equal(t, []byte(s), message.Payload)
20 | }
21 | 
22 | func TestNewBinaryMessage(t *testing.T) {
23 | 	b := []byte("TEST")
24 | 	message := wspubsub.NewBinaryMessage(b)
25 | 	require.Equal(t, wspubsub.MessageTypeBinary, message.Type)
26 | 	require.Equal(t, b, message.Payload)
27 | 
28 | 	s := "TEST"
29 | 	message = wspubsub.NewBinaryMessageFromString(s)
30 | 	require.Equal(t, wspubsub.MessageTypeBinary, message.Type)
31 | 	require.Equal(t, []byte(s), message.Payload)
32 | }
33 | 


--------------------------------------------------------------------------------
/mock/client_factory.go:
--------------------------------------------------------------------------------
  1 | // Code generated by MockGen. DO NOT EDIT.
  2 | // Source: client_factory.go
  3 | 
  4 | // Package mock is a generated GoMock package.
  5 | package mock
  6 | 
  7 | import (
  8 | 	http "net/http"
  9 | 	reflect "reflect"
 10 | 
 11 | 	gomock "github.com/golang/mock/gomock"
 12 | 	wspubsub "github.com/kpeu3i/wspubsub"
 13 | )
 14 | 
 15 | // MockWebsocketConnectionUpgrader is a mock of WebsocketConnectionUpgrader interface
 16 | type MockWebsocketConnectionUpgrader struct {
 17 | 	ctrl     *gomock.Controller
 18 | 	recorder *MockWebsocketConnectionUpgraderMockRecorder
 19 | }
 20 | 
 21 | // MockWebsocketConnectionUpgraderMockRecorder is the mock recorder for MockWebsocketConnectionUpgrader
 22 | type MockWebsocketConnectionUpgraderMockRecorder struct {
 23 | 	mock *MockWebsocketConnectionUpgrader
 24 | }
 25 | 
 26 | // NewMockWebsocketConnectionUpgrader creates a new mock instance
 27 | func NewMockWebsocketConnectionUpgrader(ctrl *gomock.Controller) *MockWebsocketConnectionUpgrader {
 28 | 	mock := &MockWebsocketConnectionUpgrader{ctrl: ctrl}
 29 | 	mock.recorder = &MockWebsocketConnectionUpgraderMockRecorder{mock}
 30 | 	return mock
 31 | }
 32 | 
 33 | // EXPECT returns an object that allows the caller to indicate expected use
 34 | func (m *MockWebsocketConnectionUpgrader) EXPECT() *MockWebsocketConnectionUpgraderMockRecorder {
 35 | 	return m.recorder
 36 | }
 37 | 
 38 | // Upgrade mocks base method
 39 | func (m *MockWebsocketConnectionUpgrader) Upgrade(w http.ResponseWriter, r *http.Request) (wspubsub.WebsocketConnection, error) {
 40 | 	m.ctrl.T.Helper()
 41 | 	ret := m.ctrl.Call(m, "Upgrade", w, r)
 42 | 	ret0, _ := ret[0].(wspubsub.WebsocketConnection)
 43 | 	ret1, _ := ret[1].(error)
 44 | 	return ret0, ret1
 45 | }
 46 | 
 47 | // Upgrade indicates an expected call of Upgrade
 48 | func (mr *MockWebsocketConnectionUpgraderMockRecorder) Upgrade(w, r interface{}) *gomock.Call {
 49 | 	mr.mock.ctrl.T.Helper()
 50 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Upgrade", reflect.TypeOf((*MockWebsocketConnectionUpgrader)(nil).Upgrade), w, r)
 51 | }
 52 | 
 53 | // MockWebsocketConnection is a mock of WebsocketConnection interface
 54 | type MockWebsocketConnection struct {
 55 | 	ctrl     *gomock.Controller
 56 | 	recorder *MockWebsocketConnectionMockRecorder
 57 | }
 58 | 
 59 | // MockWebsocketConnectionMockRecorder is the mock recorder for MockWebsocketConnection
 60 | type MockWebsocketConnectionMockRecorder struct {
 61 | 	mock *MockWebsocketConnection
 62 | }
 63 | 
 64 | // NewMockWebsocketConnection creates a new mock instance
 65 | func NewMockWebsocketConnection(ctrl *gomock.Controller) *MockWebsocketConnection {
 66 | 	mock := &MockWebsocketConnection{ctrl: ctrl}
 67 | 	mock.recorder = &MockWebsocketConnectionMockRecorder{mock}
 68 | 	return mock
 69 | }
 70 | 
 71 | // EXPECT returns an object that allows the caller to indicate expected use
 72 | func (m *MockWebsocketConnection) EXPECT() *MockWebsocketConnectionMockRecorder {
 73 | 	return m.recorder
 74 | }
 75 | 
 76 | // Read mocks base method
 77 | func (m *MockWebsocketConnection) Read() (wspubsub.Message, error) {
 78 | 	m.ctrl.T.Helper()
 79 | 	ret := m.ctrl.Call(m, "Read")
 80 | 	ret0, _ := ret[0].(wspubsub.Message)
 81 | 	ret1, _ := ret[1].(error)
 82 | 	return ret0, ret1
 83 | }
 84 | 
 85 | // Read indicates an expected call of Read
 86 | func (mr *MockWebsocketConnectionMockRecorder) Read() *gomock.Call {
 87 | 	mr.mock.ctrl.T.Helper()
 88 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Read", reflect.TypeOf((*MockWebsocketConnection)(nil).Read))
 89 | }
 90 | 
 91 | // Write mocks base method
 92 | func (m *MockWebsocketConnection) Write(message wspubsub.Message) error {
 93 | 	m.ctrl.T.Helper()
 94 | 	ret := m.ctrl.Call(m, "Write", message)
 95 | 	ret0, _ := ret[0].(error)
 96 | 	return ret0
 97 | }
 98 | 
 99 | // Write indicates an expected call of Write
100 | func (mr *MockWebsocketConnectionMockRecorder) Write(message interface{}) *gomock.Call {
101 | 	mr.mock.ctrl.T.Helper()
102 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Write", reflect.TypeOf((*MockWebsocketConnection)(nil).Write), message)
103 | }
104 | 
105 | // Close mocks base method
106 | func (m *MockWebsocketConnection) Close() error {
107 | 	m.ctrl.T.Helper()
108 | 	ret := m.ctrl.Call(m, "Close")
109 | 	ret0, _ := ret[0].(error)
110 | 	return ret0
111 | }
112 | 
113 | // Close indicates an expected call of Close
114 | func (mr *MockWebsocketConnectionMockRecorder) Close() *gomock.Call {
115 | 	mr.mock.ctrl.T.Helper()
116 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockWebsocketConnection)(nil).Close))
117 | }
118 | 
119 | // MockUUIDGenerator is a mock of UUIDGenerator interface
120 | type MockUUIDGenerator struct {
121 | 	ctrl     *gomock.Controller
122 | 	recorder *MockUUIDGeneratorMockRecorder
123 | }
124 | 
125 | // MockUUIDGeneratorMockRecorder is the mock recorder for MockUUIDGenerator
126 | type MockUUIDGeneratorMockRecorder struct {
127 | 	mock *MockUUIDGenerator
128 | }
129 | 
130 | // NewMockUUIDGenerator creates a new mock instance
131 | func NewMockUUIDGenerator(ctrl *gomock.Controller) *MockUUIDGenerator {
132 | 	mock := &MockUUIDGenerator{ctrl: ctrl}
133 | 	mock.recorder = &MockUUIDGeneratorMockRecorder{mock}
134 | 	return mock
135 | }
136 | 
137 | // EXPECT returns an object that allows the caller to indicate expected use
138 | func (m *MockUUIDGenerator) EXPECT() *MockUUIDGeneratorMockRecorder {
139 | 	return m.recorder
140 | }
141 | 
142 | // GenerateV4 mocks base method
143 | func (m *MockUUIDGenerator) GenerateV4() wspubsub.UUID {
144 | 	m.ctrl.T.Helper()
145 | 	ret := m.ctrl.Call(m, "GenerateV4")
146 | 	ret0, _ := ret[0].(wspubsub.UUID)
147 | 	return ret0
148 | }
149 | 
150 | // GenerateV4 indicates an expected call of GenerateV4
151 | func (mr *MockUUIDGeneratorMockRecorder) GenerateV4() *gomock.Call {
152 | 	mr.mock.ctrl.T.Helper()
153 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GenerateV4", reflect.TypeOf((*MockUUIDGenerator)(nil).GenerateV4))
154 | }
155 | 


--------------------------------------------------------------------------------
/mock/hub.go:
--------------------------------------------------------------------------------
  1 | // Code generated by MockGen. DO NOT EDIT.
  2 | // Source: hub.go
  3 | 
  4 | // Package mock is a generated GoMock package.
  5 | package mock
  6 | 
  7 | import (
  8 | 	http "net/http"
  9 | 	reflect "reflect"
 10 | 
 11 | 	gomock "github.com/golang/mock/gomock"
 12 | 	wspubsub "github.com/kpeu3i/wspubsub"
 13 | )
 14 | 
 15 | // MockWebsocketClient is a mock of WebsocketClient interface
 16 | type MockWebsocketClient struct {
 17 | 	ctrl     *gomock.Controller
 18 | 	recorder *MockWebsocketClientMockRecorder
 19 | }
 20 | 
 21 | // MockWebsocketClientMockRecorder is the mock recorder for MockWebsocketClient
 22 | type MockWebsocketClientMockRecorder struct {
 23 | 	mock *MockWebsocketClient
 24 | }
 25 | 
 26 | // NewMockWebsocketClient creates a new mock instance
 27 | func NewMockWebsocketClient(ctrl *gomock.Controller) *MockWebsocketClient {
 28 | 	mock := &MockWebsocketClient{ctrl: ctrl}
 29 | 	mock.recorder = &MockWebsocketClientMockRecorder{mock}
 30 | 	return mock
 31 | }
 32 | 
 33 | // EXPECT returns an object that allows the caller to indicate expected use
 34 | func (m *MockWebsocketClient) EXPECT() *MockWebsocketClientMockRecorder {
 35 | 	return m.recorder
 36 | }
 37 | 
 38 | // ID mocks base method
 39 | func (m *MockWebsocketClient) ID() wspubsub.UUID {
 40 | 	m.ctrl.T.Helper()
 41 | 	ret := m.ctrl.Call(m, "ID")
 42 | 	ret0, _ := ret[0].(wspubsub.UUID)
 43 | 	return ret0
 44 | }
 45 | 
 46 | // ID indicates an expected call of ID
 47 | func (mr *MockWebsocketClientMockRecorder) ID() *gomock.Call {
 48 | 	mr.mock.ctrl.T.Helper()
 49 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ID", reflect.TypeOf((*MockWebsocketClient)(nil).ID))
 50 | }
 51 | 
 52 | // Connect mocks base method
 53 | func (m *MockWebsocketClient) Connect(response http.ResponseWriter, request *http.Request) error {
 54 | 	m.ctrl.T.Helper()
 55 | 	ret := m.ctrl.Call(m, "Connect", response, request)
 56 | 	ret0, _ := ret[0].(error)
 57 | 	return ret0
 58 | }
 59 | 
 60 | // Connect indicates an expected call of Connect
 61 | func (mr *MockWebsocketClientMockRecorder) Connect(response, request interface{}) *gomock.Call {
 62 | 	mr.mock.ctrl.T.Helper()
 63 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Connect", reflect.TypeOf((*MockWebsocketClient)(nil).Connect), response, request)
 64 | }
 65 | 
 66 | // OnReceive mocks base method
 67 | func (m *MockWebsocketClient) OnReceive(handler wspubsub.ReceiveHandler) {
 68 | 	m.ctrl.T.Helper()
 69 | 	m.ctrl.Call(m, "OnReceive", handler)
 70 | }
 71 | 
 72 | // OnReceive indicates an expected call of OnReceive
 73 | func (mr *MockWebsocketClientMockRecorder) OnReceive(handler interface{}) *gomock.Call {
 74 | 	mr.mock.ctrl.T.Helper()
 75 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnReceive", reflect.TypeOf((*MockWebsocketClient)(nil).OnReceive), handler)
 76 | }
 77 | 
 78 | // OnError mocks base method
 79 | func (m *MockWebsocketClient) OnError(handler wspubsub.ErrorHandler) {
 80 | 	m.ctrl.T.Helper()
 81 | 	m.ctrl.Call(m, "OnError", handler)
 82 | }
 83 | 
 84 | // OnError indicates an expected call of OnError
 85 | func (mr *MockWebsocketClientMockRecorder) OnError(handler interface{}) *gomock.Call {
 86 | 	mr.mock.ctrl.T.Helper()
 87 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnError", reflect.TypeOf((*MockWebsocketClient)(nil).OnError), handler)
 88 | }
 89 | 
 90 | // Send mocks base method
 91 | func (m *MockWebsocketClient) Send(message wspubsub.Message) error {
 92 | 	m.ctrl.T.Helper()
 93 | 	ret := m.ctrl.Call(m, "Send", message)
 94 | 	ret0, _ := ret[0].(error)
 95 | 	return ret0
 96 | }
 97 | 
 98 | // Send indicates an expected call of Send
 99 | func (mr *MockWebsocketClientMockRecorder) Send(message interface{}) *gomock.Call {
100 | 	mr.mock.ctrl.T.Helper()
101 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Send", reflect.TypeOf((*MockWebsocketClient)(nil).Send), message)
102 | }
103 | 
104 | // Close mocks base method
105 | func (m *MockWebsocketClient) Close() error {
106 | 	m.ctrl.T.Helper()
107 | 	ret := m.ctrl.Call(m, "Close")
108 | 	ret0, _ := ret[0].(error)
109 | 	return ret0
110 | }
111 | 
112 | // Close indicates an expected call of Close
113 | func (mr *MockWebsocketClientMockRecorder) Close() *gomock.Call {
114 | 	mr.mock.ctrl.T.Helper()
115 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockWebsocketClient)(nil).Close))
116 | }
117 | 
118 | // MockWebsocketClientStore is a mock of WebsocketClientStore interface
119 | type MockWebsocketClientStore struct {
120 | 	ctrl     *gomock.Controller
121 | 	recorder *MockWebsocketClientStoreMockRecorder
122 | }
123 | 
124 | // MockWebsocketClientStoreMockRecorder is the mock recorder for MockWebsocketClientStore
125 | type MockWebsocketClientStoreMockRecorder struct {
126 | 	mock *MockWebsocketClientStore
127 | }
128 | 
129 | // NewMockWebsocketClientStore creates a new mock instance
130 | func NewMockWebsocketClientStore(ctrl *gomock.Controller) *MockWebsocketClientStore {
131 | 	mock := &MockWebsocketClientStore{ctrl: ctrl}
132 | 	mock.recorder = &MockWebsocketClientStoreMockRecorder{mock}
133 | 	return mock
134 | }
135 | 
136 | // EXPECT returns an object that allows the caller to indicate expected use
137 | func (m *MockWebsocketClientStore) EXPECT() *MockWebsocketClientStoreMockRecorder {
138 | 	return m.recorder
139 | }
140 | 
141 | // Get mocks base method
142 | func (m *MockWebsocketClientStore) Get(clientID wspubsub.UUID) (wspubsub.WebsocketClient, error) {
143 | 	m.ctrl.T.Helper()
144 | 	ret := m.ctrl.Call(m, "Get", clientID)
145 | 	ret0, _ := ret[0].(wspubsub.WebsocketClient)
146 | 	ret1, _ := ret[1].(error)
147 | 	return ret0, ret1
148 | }
149 | 
150 | // Get indicates an expected call of Get
151 | func (mr *MockWebsocketClientStoreMockRecorder) Get(clientID interface{}) *gomock.Call {
152 | 	mr.mock.ctrl.T.Helper()
153 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get", reflect.TypeOf((*MockWebsocketClientStore)(nil).Get), clientID)
154 | }
155 | 
156 | // Set mocks base method
157 | func (m *MockWebsocketClientStore) Set(client wspubsub.WebsocketClient) {
158 | 	m.ctrl.T.Helper()
159 | 	m.ctrl.Call(m, "Set", client)
160 | }
161 | 
162 | // Set indicates an expected call of Set
163 | func (mr *MockWebsocketClientStoreMockRecorder) Set(client interface{}) *gomock.Call {
164 | 	mr.mock.ctrl.T.Helper()
165 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Set", reflect.TypeOf((*MockWebsocketClientStore)(nil).Set), client)
166 | }
167 | 
168 | // Unset mocks base method
169 | func (m *MockWebsocketClientStore) Unset(clientID wspubsub.UUID) error {
170 | 	m.ctrl.T.Helper()
171 | 	ret := m.ctrl.Call(m, "Unset", clientID)
172 | 	ret0, _ := ret[0].(error)
173 | 	return ret0
174 | }
175 | 
176 | // Unset indicates an expected call of Unset
177 | func (mr *MockWebsocketClientStoreMockRecorder) Unset(clientID interface{}) *gomock.Call {
178 | 	mr.mock.ctrl.T.Helper()
179 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Unset", reflect.TypeOf((*MockWebsocketClientStore)(nil).Unset), clientID)
180 | }
181 | 
182 | // Count mocks base method
183 | func (m *MockWebsocketClientStore) Count(channels ...string) int {
184 | 	m.ctrl.T.Helper()
185 | 	varargs := []interface{}{}
186 | 	for _, a := range channels {
187 | 		varargs = append(varargs, a)
188 | 	}
189 | 	ret := m.ctrl.Call(m, "Count", varargs...)
190 | 	ret0, _ := ret[0].(int)
191 | 	return ret0
192 | }
193 | 
194 | // Count indicates an expected call of Count
195 | func (mr *MockWebsocketClientStoreMockRecorder) Count(channels ...interface{}) *gomock.Call {
196 | 	mr.mock.ctrl.T.Helper()
197 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Count", reflect.TypeOf((*MockWebsocketClientStore)(nil).Count), channels...)
198 | }
199 | 
200 | // Find mocks base method
201 | func (m *MockWebsocketClientStore) Find(fn wspubsub.IterateFunc, channels ...string) error {
202 | 	m.ctrl.T.Helper()
203 | 	varargs := []interface{}{fn}
204 | 	for _, a := range channels {
205 | 		varargs = append(varargs, a)
206 | 	}
207 | 	ret := m.ctrl.Call(m, "Find", varargs...)
208 | 	ret0, _ := ret[0].(error)
209 | 	return ret0
210 | }
211 | 
212 | // Find indicates an expected call of Find
213 | func (mr *MockWebsocketClientStoreMockRecorder) Find(fn interface{}, channels ...interface{}) *gomock.Call {
214 | 	mr.mock.ctrl.T.Helper()
215 | 	varargs := append([]interface{}{fn}, channels...)
216 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Find", reflect.TypeOf((*MockWebsocketClientStore)(nil).Find), varargs...)
217 | }
218 | 
219 | // Channels mocks base method
220 | func (m *MockWebsocketClientStore) Channels(clientID wspubsub.UUID) ([]string, error) {
221 | 	m.ctrl.T.Helper()
222 | 	ret := m.ctrl.Call(m, "Channels", clientID)
223 | 	ret0, _ := ret[0].([]string)
224 | 	ret1, _ := ret[1].(error)
225 | 	return ret0, ret1
226 | }
227 | 
228 | // Channels indicates an expected call of Channels
229 | func (mr *MockWebsocketClientStoreMockRecorder) Channels(clientID interface{}) *gomock.Call {
230 | 	mr.mock.ctrl.T.Helper()
231 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Channels", reflect.TypeOf((*MockWebsocketClientStore)(nil).Channels), clientID)
232 | }
233 | 
234 | // CountChannels mocks base method
235 | func (m *MockWebsocketClientStore) CountChannels(clientID wspubsub.UUID) (int, error) {
236 | 	m.ctrl.T.Helper()
237 | 	ret := m.ctrl.Call(m, "CountChannels", clientID)
238 | 	ret0, _ := ret[0].(int)
239 | 	ret1, _ := ret[1].(error)
240 | 	return ret0, ret1
241 | }
242 | 
243 | // CountChannels indicates an expected call of CountChannels
244 | func (mr *MockWebsocketClientStoreMockRecorder) CountChannels(clientID interface{}) *gomock.Call {
245 | 	mr.mock.ctrl.T.Helper()
246 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CountChannels", reflect.TypeOf((*MockWebsocketClientStore)(nil).CountChannels), clientID)
247 | }
248 | 
249 | // SetChannels mocks base method
250 | func (m *MockWebsocketClientStore) SetChannels(clientID wspubsub.UUID, channels ...string) error {
251 | 	m.ctrl.T.Helper()
252 | 	varargs := []interface{}{clientID}
253 | 	for _, a := range channels {
254 | 		varargs = append(varargs, a)
255 | 	}
256 | 	ret := m.ctrl.Call(m, "SetChannels", varargs...)
257 | 	ret0, _ := ret[0].(error)
258 | 	return ret0
259 | }
260 | 
261 | // SetChannels indicates an expected call of SetChannels
262 | func (mr *MockWebsocketClientStoreMockRecorder) SetChannels(clientID interface{}, channels ...interface{}) *gomock.Call {
263 | 	mr.mock.ctrl.T.Helper()
264 | 	varargs := append([]interface{}{clientID}, channels...)
265 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetChannels", reflect.TypeOf((*MockWebsocketClientStore)(nil).SetChannels), varargs...)
266 | }
267 | 
268 | // UnsetChannels mocks base method
269 | func (m *MockWebsocketClientStore) UnsetChannels(clientID wspubsub.UUID, channels ...string) error {
270 | 	m.ctrl.T.Helper()
271 | 	varargs := []interface{}{clientID}
272 | 	for _, a := range channels {
273 | 		varargs = append(varargs, a)
274 | 	}
275 | 	ret := m.ctrl.Call(m, "UnsetChannels", varargs...)
276 | 	ret0, _ := ret[0].(error)
277 | 	return ret0
278 | }
279 | 
280 | // UnsetChannels indicates an expected call of UnsetChannels
281 | func (mr *MockWebsocketClientStoreMockRecorder) UnsetChannels(clientID interface{}, channels ...interface{}) *gomock.Call {
282 | 	mr.mock.ctrl.T.Helper()
283 | 	varargs := append([]interface{}{clientID}, channels...)
284 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UnsetChannels", reflect.TypeOf((*MockWebsocketClientStore)(nil).UnsetChannels), varargs...)
285 | }
286 | 
287 | // MockWebsocketClientFactory is a mock of WebsocketClientFactory interface
288 | type MockWebsocketClientFactory struct {
289 | 	ctrl     *gomock.Controller
290 | 	recorder *MockWebsocketClientFactoryMockRecorder
291 | }
292 | 
293 | // MockWebsocketClientFactoryMockRecorder is the mock recorder for MockWebsocketClientFactory
294 | type MockWebsocketClientFactoryMockRecorder struct {
295 | 	mock *MockWebsocketClientFactory
296 | }
297 | 
298 | // NewMockWebsocketClientFactory creates a new mock instance
299 | func NewMockWebsocketClientFactory(ctrl *gomock.Controller) *MockWebsocketClientFactory {
300 | 	mock := &MockWebsocketClientFactory{ctrl: ctrl}
301 | 	mock.recorder = &MockWebsocketClientFactoryMockRecorder{mock}
302 | 	return mock
303 | }
304 | 
305 | // EXPECT returns an object that allows the caller to indicate expected use
306 | func (m *MockWebsocketClientFactory) EXPECT() *MockWebsocketClientFactoryMockRecorder {
307 | 	return m.recorder
308 | }
309 | 
310 | // Create mocks base method
311 | func (m *MockWebsocketClientFactory) Create() wspubsub.WebsocketClient {
312 | 	m.ctrl.T.Helper()
313 | 	ret := m.ctrl.Call(m, "Create")
314 | 	ret0, _ := ret[0].(wspubsub.WebsocketClient)
315 | 	return ret0
316 | }
317 | 
318 | // Create indicates an expected call of Create
319 | func (mr *MockWebsocketClientFactoryMockRecorder) Create() *gomock.Call {
320 | 	mr.mock.ctrl.T.Helper()
321 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Create", reflect.TypeOf((*MockWebsocketClientFactory)(nil).Create))
322 | }
323 | 
324 | // MockLogger is a mock of Logger interface
325 | type MockLogger struct {
326 | 	ctrl     *gomock.Controller
327 | 	recorder *MockLoggerMockRecorder
328 | }
329 | 
330 | // MockLoggerMockRecorder is the mock recorder for MockLogger
331 | type MockLoggerMockRecorder struct {
332 | 	mock *MockLogger
333 | }
334 | 
335 | // NewMockLogger creates a new mock instance
336 | func NewMockLogger(ctrl *gomock.Controller) *MockLogger {
337 | 	mock := &MockLogger{ctrl: ctrl}
338 | 	mock.recorder = &MockLoggerMockRecorder{mock}
339 | 	return mock
340 | }
341 | 
342 | // EXPECT returns an object that allows the caller to indicate expected use
343 | func (m *MockLogger) EXPECT() *MockLoggerMockRecorder {
344 | 	return m.recorder
345 | }
346 | 
347 | // Debug mocks base method
348 | func (m *MockLogger) Debug(args ...interface{}) {
349 | 	m.ctrl.T.Helper()
350 | 	varargs := []interface{}{}
351 | 	for _, a := range args {
352 | 		varargs = append(varargs, a)
353 | 	}
354 | 	m.ctrl.Call(m, "Debug", varargs...)
355 | }
356 | 
357 | // Debug indicates an expected call of Debug
358 | func (mr *MockLoggerMockRecorder) Debug(args ...interface{}) *gomock.Call {
359 | 	mr.mock.ctrl.T.Helper()
360 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Debug", reflect.TypeOf((*MockLogger)(nil).Debug), args...)
361 | }
362 | 
363 | // Info mocks base method
364 | func (m *MockLogger) Info(args ...interface{}) {
365 | 	m.ctrl.T.Helper()
366 | 	varargs := []interface{}{}
367 | 	for _, a := range args {
368 | 		varargs = append(varargs, a)
369 | 	}
370 | 	m.ctrl.Call(m, "Info", varargs...)
371 | }
372 | 
373 | // Info indicates an expected call of Info
374 | func (mr *MockLoggerMockRecorder) Info(args ...interface{}) *gomock.Call {
375 | 	mr.mock.ctrl.T.Helper()
376 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Info", reflect.TypeOf((*MockLogger)(nil).Info), args...)
377 | }
378 | 
379 | // Print mocks base method
380 | func (m *MockLogger) Print(args ...interface{}) {
381 | 	m.ctrl.T.Helper()
382 | 	varargs := []interface{}{}
383 | 	for _, a := range args {
384 | 		varargs = append(varargs, a)
385 | 	}
386 | 	m.ctrl.Call(m, "Print", varargs...)
387 | }
388 | 
389 | // Print indicates an expected call of Print
390 | func (mr *MockLoggerMockRecorder) Print(args ...interface{}) *gomock.Call {
391 | 	mr.mock.ctrl.T.Helper()
392 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Print", reflect.TypeOf((*MockLogger)(nil).Print), args...)
393 | }
394 | 
395 | // Warn mocks base method
396 | func (m *MockLogger) Warn(args ...interface{}) {
397 | 	m.ctrl.T.Helper()
398 | 	varargs := []interface{}{}
399 | 	for _, a := range args {
400 | 		varargs = append(varargs, a)
401 | 	}
402 | 	m.ctrl.Call(m, "Warn", varargs...)
403 | }
404 | 
405 | // Warn indicates an expected call of Warn
406 | func (mr *MockLoggerMockRecorder) Warn(args ...interface{}) *gomock.Call {
407 | 	mr.mock.ctrl.T.Helper()
408 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Warn", reflect.TypeOf((*MockLogger)(nil).Warn), args...)
409 | }
410 | 
411 | // Error mocks base method
412 | func (m *MockLogger) Error(args ...interface{}) {
413 | 	m.ctrl.T.Helper()
414 | 	varargs := []interface{}{}
415 | 	for _, a := range args {
416 | 		varargs = append(varargs, a)
417 | 	}
418 | 	m.ctrl.Call(m, "Error", varargs...)
419 | }
420 | 
421 | // Error indicates an expected call of Error
422 | func (mr *MockLoggerMockRecorder) Error(args ...interface{}) *gomock.Call {
423 | 	mr.mock.ctrl.T.Helper()
424 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Error", reflect.TypeOf((*MockLogger)(nil).Error), args...)
425 | }
426 | 
427 | // Fatal mocks base method
428 | func (m *MockLogger) Fatal(args ...interface{}) {
429 | 	m.ctrl.T.Helper()
430 | 	varargs := []interface{}{}
431 | 	for _, a := range args {
432 | 		varargs = append(varargs, a)
433 | 	}
434 | 	m.ctrl.Call(m, "Fatal", varargs...)
435 | }
436 | 
437 | // Fatal indicates an expected call of Fatal
438 | func (mr *MockLoggerMockRecorder) Fatal(args ...interface{}) *gomock.Call {
439 | 	mr.mock.ctrl.T.Helper()
440 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Fatal", reflect.TypeOf((*MockLogger)(nil).Fatal), args...)
441 | }
442 | 
443 | // Panic mocks base method
444 | func (m *MockLogger) Panic(args ...interface{}) {
445 | 	m.ctrl.T.Helper()
446 | 	varargs := []interface{}{}
447 | 	for _, a := range args {
448 | 		varargs = append(varargs, a)
449 | 	}
450 | 	m.ctrl.Call(m, "Panic", varargs...)
451 | }
452 | 
453 | // Panic indicates an expected call of Panic
454 | func (mr *MockLoggerMockRecorder) Panic(args ...interface{}) *gomock.Call {
455 | 	mr.mock.ctrl.T.Helper()
456 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Panic", reflect.TypeOf((*MockLogger)(nil).Panic), args...)
457 | }
458 | 
459 | // Debugln mocks base method
460 | func (m *MockLogger) Debugln(args ...interface{}) {
461 | 	m.ctrl.T.Helper()
462 | 	varargs := []interface{}{}
463 | 	for _, a := range args {
464 | 		varargs = append(varargs, a)
465 | 	}
466 | 	m.ctrl.Call(m, "Debugln", varargs...)
467 | }
468 | 
469 | // Debugln indicates an expected call of Debugln
470 | func (mr *MockLoggerMockRecorder) Debugln(args ...interface{}) *gomock.Call {
471 | 	mr.mock.ctrl.T.Helper()
472 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Debugln", reflect.TypeOf((*MockLogger)(nil).Debugln), args...)
473 | }
474 | 
475 | // Infoln mocks base method
476 | func (m *MockLogger) Infoln(args ...interface{}) {
477 | 	m.ctrl.T.Helper()
478 | 	varargs := []interface{}{}
479 | 	for _, a := range args {
480 | 		varargs = append(varargs, a)
481 | 	}
482 | 	m.ctrl.Call(m, "Infoln", varargs...)
483 | }
484 | 
485 | // Infoln indicates an expected call of Infoln
486 | func (mr *MockLoggerMockRecorder) Infoln(args ...interface{}) *gomock.Call {
487 | 	mr.mock.ctrl.T.Helper()
488 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Infoln", reflect.TypeOf((*MockLogger)(nil).Infoln), args...)
489 | }
490 | 
491 | // Println mocks base method
492 | func (m *MockLogger) Println(args ...interface{}) {
493 | 	m.ctrl.T.Helper()
494 | 	varargs := []interface{}{}
495 | 	for _, a := range args {
496 | 		varargs = append(varargs, a)
497 | 	}
498 | 	m.ctrl.Call(m, "Println", varargs...)
499 | }
500 | 
501 | // Println indicates an expected call of Println
502 | func (mr *MockLoggerMockRecorder) Println(args ...interface{}) *gomock.Call {
503 | 	mr.mock.ctrl.T.Helper()
504 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Println", reflect.TypeOf((*MockLogger)(nil).Println), args...)
505 | }
506 | 
507 | // Warnln mocks base method
508 | func (m *MockLogger) Warnln(args ...interface{}) {
509 | 	m.ctrl.T.Helper()
510 | 	varargs := []interface{}{}
511 | 	for _, a := range args {
512 | 		varargs = append(varargs, a)
513 | 	}
514 | 	m.ctrl.Call(m, "Warnln", varargs...)
515 | }
516 | 
517 | // Warnln indicates an expected call of Warnln
518 | func (mr *MockLoggerMockRecorder) Warnln(args ...interface{}) *gomock.Call {
519 | 	mr.mock.ctrl.T.Helper()
520 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Warnln", reflect.TypeOf((*MockLogger)(nil).Warnln), args...)
521 | }
522 | 
523 | // Errorln mocks base method
524 | func (m *MockLogger) Errorln(args ...interface{}) {
525 | 	m.ctrl.T.Helper()
526 | 	varargs := []interface{}{}
527 | 	for _, a := range args {
528 | 		varargs = append(varargs, a)
529 | 	}
530 | 	m.ctrl.Call(m, "Errorln", varargs...)
531 | }
532 | 
533 | // Errorln indicates an expected call of Errorln
534 | func (mr *MockLoggerMockRecorder) Errorln(args ...interface{}) *gomock.Call {
535 | 	mr.mock.ctrl.T.Helper()
536 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Errorln", reflect.TypeOf((*MockLogger)(nil).Errorln), args...)
537 | }
538 | 
539 | // Fatalln mocks base method
540 | func (m *MockLogger) Fatalln(args ...interface{}) {
541 | 	m.ctrl.T.Helper()
542 | 	varargs := []interface{}{}
543 | 	for _, a := range args {
544 | 		varargs = append(varargs, a)
545 | 	}
546 | 	m.ctrl.Call(m, "Fatalln", varargs...)
547 | }
548 | 
549 | // Fatalln indicates an expected call of Fatalln
550 | func (mr *MockLoggerMockRecorder) Fatalln(args ...interface{}) *gomock.Call {
551 | 	mr.mock.ctrl.T.Helper()
552 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Fatalln", reflect.TypeOf((*MockLogger)(nil).Fatalln), args...)
553 | }
554 | 
555 | // Panicln mocks base method
556 | func (m *MockLogger) Panicln(args ...interface{}) {
557 | 	m.ctrl.T.Helper()
558 | 	varargs := []interface{}{}
559 | 	for _, a := range args {
560 | 		varargs = append(varargs, a)
561 | 	}
562 | 	m.ctrl.Call(m, "Panicln", varargs...)
563 | }
564 | 
565 | // Panicln indicates an expected call of Panicln
566 | func (mr *MockLoggerMockRecorder) Panicln(args ...interface{}) *gomock.Call {
567 | 	mr.mock.ctrl.T.Helper()
568 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Panicln", reflect.TypeOf((*MockLogger)(nil).Panicln), args...)
569 | }
570 | 
571 | // Debugf mocks base method
572 | func (m *MockLogger) Debugf(format string, args ...interface{}) {
573 | 	m.ctrl.T.Helper()
574 | 	varargs := []interface{}{format}
575 | 	for _, a := range args {
576 | 		varargs = append(varargs, a)
577 | 	}
578 | 	m.ctrl.Call(m, "Debugf", varargs...)
579 | }
580 | 
581 | // Debugf indicates an expected call of Debugf
582 | func (mr *MockLoggerMockRecorder) Debugf(format interface{}, args ...interface{}) *gomock.Call {
583 | 	mr.mock.ctrl.T.Helper()
584 | 	varargs := append([]interface{}{format}, args...)
585 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Debugf", reflect.TypeOf((*MockLogger)(nil).Debugf), varargs...)
586 | }
587 | 
588 | // Infof mocks base method
589 | func (m *MockLogger) Infof(format string, args ...interface{}) {
590 | 	m.ctrl.T.Helper()
591 | 	varargs := []interface{}{format}
592 | 	for _, a := range args {
593 | 		varargs = append(varargs, a)
594 | 	}
595 | 	m.ctrl.Call(m, "Infof", varargs...)
596 | }
597 | 
598 | // Infof indicates an expected call of Infof
599 | func (mr *MockLoggerMockRecorder) Infof(format interface{}, args ...interface{}) *gomock.Call {
600 | 	mr.mock.ctrl.T.Helper()
601 | 	varargs := append([]interface{}{format}, args...)
602 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Infof", reflect.TypeOf((*MockLogger)(nil).Infof), varargs...)
603 | }
604 | 
605 | // Printf mocks base method
606 | func (m *MockLogger) Printf(format string, args ...interface{}) {
607 | 	m.ctrl.T.Helper()
608 | 	varargs := []interface{}{format}
609 | 	for _, a := range args {
610 | 		varargs = append(varargs, a)
611 | 	}
612 | 	m.ctrl.Call(m, "Printf", varargs...)
613 | }
614 | 
615 | // Printf indicates an expected call of Printf
616 | func (mr *MockLoggerMockRecorder) Printf(format interface{}, args ...interface{}) *gomock.Call {
617 | 	mr.mock.ctrl.T.Helper()
618 | 	varargs := append([]interface{}{format}, args...)
619 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Printf", reflect.TypeOf((*MockLogger)(nil).Printf), varargs...)
620 | }
621 | 
622 | // Warnf mocks base method
623 | func (m *MockLogger) Warnf(format string, args ...interface{}) {
624 | 	m.ctrl.T.Helper()
625 | 	varargs := []interface{}{format}
626 | 	for _, a := range args {
627 | 		varargs = append(varargs, a)
628 | 	}
629 | 	m.ctrl.Call(m, "Warnf", varargs...)
630 | }
631 | 
632 | // Warnf indicates an expected call of Warnf
633 | func (mr *MockLoggerMockRecorder) Warnf(format interface{}, args ...interface{}) *gomock.Call {
634 | 	mr.mock.ctrl.T.Helper()
635 | 	varargs := append([]interface{}{format}, args...)
636 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Warnf", reflect.TypeOf((*MockLogger)(nil).Warnf), varargs...)
637 | }
638 | 
639 | // Errorf mocks base method
640 | func (m *MockLogger) Errorf(format string, args ...interface{}) {
641 | 	m.ctrl.T.Helper()
642 | 	varargs := []interface{}{format}
643 | 	for _, a := range args {
644 | 		varargs = append(varargs, a)
645 | 	}
646 | 	m.ctrl.Call(m, "Errorf", varargs...)
647 | }
648 | 
649 | // Errorf indicates an expected call of Errorf
650 | func (mr *MockLoggerMockRecorder) Errorf(format interface{}, args ...interface{}) *gomock.Call {
651 | 	mr.mock.ctrl.T.Helper()
652 | 	varargs := append([]interface{}{format}, args...)
653 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Errorf", reflect.TypeOf((*MockLogger)(nil).Errorf), varargs...)
654 | }
655 | 
656 | // Fatalf mocks base method
657 | func (m *MockLogger) Fatalf(format string, args ...interface{}) {
658 | 	m.ctrl.T.Helper()
659 | 	varargs := []interface{}{format}
660 | 	for _, a := range args {
661 | 		varargs = append(varargs, a)
662 | 	}
663 | 	m.ctrl.Call(m, "Fatalf", varargs...)
664 | }
665 | 
666 | // Fatalf indicates an expected call of Fatalf
667 | func (mr *MockLoggerMockRecorder) Fatalf(format interface{}, args ...interface{}) *gomock.Call {
668 | 	mr.mock.ctrl.T.Helper()
669 | 	varargs := append([]interface{}{format}, args...)
670 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Fatalf", reflect.TypeOf((*MockLogger)(nil).Fatalf), varargs...)
671 | }
672 | 
673 | // Panicf mocks base method
674 | func (m *MockLogger) Panicf(format string, args ...interface{}) {
675 | 	m.ctrl.T.Helper()
676 | 	varargs := []interface{}{format}
677 | 	for _, a := range args {
678 | 		varargs = append(varargs, a)
679 | 	}
680 | 	m.ctrl.Call(m, "Panicf", varargs...)
681 | }
682 | 
683 | // Panicf indicates an expected call of Panicf
684 | func (mr *MockLoggerMockRecorder) Panicf(format interface{}, args ...interface{}) *gomock.Call {
685 | 	mr.mock.ctrl.T.Helper()
686 | 	varargs := append([]interface{}{format}, args...)
687 | 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Panicf", reflect.TypeOf((*MockLogger)(nil).Panicf), varargs...)
688 | }
689 | 


--------------------------------------------------------------------------------
/satori_uuid_generator.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import uuid "github.com/satori/go.uuid"
 4 | 
 5 | var _ UUIDGenerator = (*SatoriUUIDGenerator)(nil)
 6 | 
 7 | // SatoriUUIDGenerator is an implementation of UUIDGenerator.
 8 | type SatoriUUIDGenerator struct{}
 9 | 
10 | // GenerateV4 returns random generated UUID.
11 | func (s SatoriUUIDGenerator) GenerateV4() UUID {
12 | 	return UUID(uuid.NewV4())
13 | }
14 | 


--------------------------------------------------------------------------------
/uuid.go:
--------------------------------------------------------------------------------
 1 | package wspubsub
 2 | 
 3 | import "encoding/hex"
 4 | 
 5 | // UUID represents a type compliant with specification described in RFC 4122.
 6 | type UUID [16]byte
 7 | 
 8 | // String returns canonical string representation of UUID.
 9 | func (u UUID) String() string {
10 | 	buf := make([]byte, 36)
11 | 
12 | 	hex.Encode(buf[0:8], u[0:4])
13 | 	buf[8] = '-'
14 | 	hex.Encode(buf[9:13], u[4:6])
15 | 	buf[13] = '-'
16 | 	hex.Encode(buf[14:18], u[6:8])
17 | 	buf[18] = '-'
18 | 	hex.Encode(buf[19:23], u[8:10])
19 | 	buf[23] = '-'
20 | 	hex.Encode(buf[24:], u[10:])
21 | 
22 | 	return string(buf)
23 | }
24 | 
25 | // Bytes returns bytes slice representation of UUID.
26 | func (u UUID) Bytes() []byte {
27 | 	return u[:]
28 | }
29 | 


--------------------------------------------------------------------------------
/uuid_test.go:
--------------------------------------------------------------------------------
 1 | package wspubsub_test
 2 | 
 3 | import (
 4 | 	"testing"
 5 | 
 6 | 	"github.com/stretchr/testify/require"
 7 | )
 8 | 
 9 | func TestUUID_Bytes(t *testing.T) {
10 | 	require.Equal(t, clientID[:], clientID.Bytes())
11 | }
12 | 
13 | func TestUUID_String(t *testing.T) {
14 | 	require.Equal(t, "01020304-0506-0708-090a-0b0c0d0e0f10", clientID.String())
15 | }
16 | 


--------------------------------------------------------------------------------