272 lines
6.5 KiB
Go
272 lines
6.5 KiB
Go
package hub
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"git.yuntuo.hk/k2safe/go-omni/ws"
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
type testConnection struct {
|
|
id string
|
|
done chan struct{}
|
|
once sync.Once
|
|
full bool
|
|
sent atomic.Int64
|
|
}
|
|
|
|
func (c *testConnection) ID() string { return c.id }
|
|
func (c *testConnection) Done() <-chan struct{} { return c.done }
|
|
func (c *testConnection) TrySend(ws.Message) error {
|
|
select {
|
|
case <-c.done:
|
|
return ws.ErrClosed
|
|
default:
|
|
}
|
|
if c.full {
|
|
return ws.ErrQueueFull
|
|
}
|
|
c.sent.Add(1)
|
|
return nil
|
|
}
|
|
func (c *testConnection) Close(int, string) error { c.once.Do(func() { close(c.done) }); return nil }
|
|
func newConnection(id string) *testConnection {
|
|
return &testConnection{id: id, done: make(chan struct{})}
|
|
}
|
|
func testHub(t *testing.T, cfg Config) *Hub {
|
|
t.Helper()
|
|
h, err := New(cfg)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
t.Cleanup(func() { _ = h.Close() })
|
|
return h
|
|
}
|
|
|
|
func TestHubRoutingBackpressureAndCleanup(t *testing.T) {
|
|
h := testHub(t, Config{})
|
|
a, b, c := newConnection("a"), newConnection("b"), newConnection("c")
|
|
b.full = true
|
|
for _, connection := range []*testConnection{a, b, c} {
|
|
if err := h.Add(connection); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
for _, id := range []string{"a", "b"} {
|
|
if err := h.Join(id, "group"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
message := ws.Message{Type: ws.Text, Data: []byte("hello")}
|
|
result, err := h.SendGroup("group", message)
|
|
if err != nil || result.Attempted != 2 || result.Queued != 1 || !errors.Is(result.Errors["b"], ws.ErrQueueFull) {
|
|
t.Fatalf("group: %#v %v", result, err)
|
|
}
|
|
if c.sent.Load() != 0 {
|
|
t.Fatal("group leaked")
|
|
}
|
|
result, err = h.Broadcast(message)
|
|
if err != nil || result.Attempted != 3 || result.Queued != 2 {
|
|
t.Fatalf("broadcast: %#v %v", result, err)
|
|
}
|
|
if err = h.Kick("a", 1000, ""); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
deadline := time.After(time.Second)
|
|
for h.Stats().Connections != 2 {
|
|
select {
|
|
case <-deadline:
|
|
t.Fatal("auto cleanup timeout")
|
|
default:
|
|
time.Sleep(time.Millisecond)
|
|
}
|
|
}
|
|
if got := h.Stats(); got.Memberships != 1 || got.Groups != 1 {
|
|
t.Fatalf("stats: %#v", got)
|
|
}
|
|
if !h.Remove("b") {
|
|
t.Fatal("remove")
|
|
}
|
|
if got := h.Stats(); got.Memberships != 0 || got.Groups != 0 {
|
|
t.Fatalf("group leak: %#v", got)
|
|
}
|
|
if _, err = h.Send("missing", message); !errors.Is(err, ErrNotFound) {
|
|
t.Fatal(err)
|
|
}
|
|
_ = h.Close()
|
|
if _, err = h.Broadcast(message); !errors.Is(err, ErrClosed) {
|
|
t.Fatal(err)
|
|
}
|
|
select {
|
|
case <-c.done:
|
|
t.Fatal("hub closed borrowed connection")
|
|
default:
|
|
}
|
|
}
|
|
|
|
func TestHubLimitsAndReusedID(t *testing.T) {
|
|
h := testHub(t, Config{MaxConnections: 2, MaxGroups: 1, MaxMemberships: 1, MaxGroupsPerConnection: 1})
|
|
first := newConnection("a")
|
|
if err := h.Add(first); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := h.Add(first); !errors.Is(err, ErrDuplicate) {
|
|
t.Fatal(err)
|
|
}
|
|
if err := h.Add(newConnection("b")); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := h.Add(newConnection("c")); !errors.Is(err, ErrLimit) {
|
|
t.Fatal(err)
|
|
}
|
|
if err := h.Join("a", "one"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := h.Join("a", "one"); err != nil {
|
|
t.Fatal("idempotent join", err)
|
|
}
|
|
for _, pair := range [][2]string{{"a", "two"}, {"b", "one"}} {
|
|
if err := h.Join(pair[0], pair[1]); !errors.Is(err, ErrLimit) {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
h.Remove("a")
|
|
second := newConnection("a")
|
|
if err := h.Add(second); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
_ = first.Close(1000, "")
|
|
if err := h.Join("a", "replacement"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := h.Send("a", ws.Message{Type: ws.Text}); err != nil {
|
|
t.Fatal("old watcher removed replacement", err)
|
|
}
|
|
var nilConnection *testConnection
|
|
if err := h.Add(nilConnection); !errors.Is(err, ErrInvalid) {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
func TestHubConcurrentLifecycle(t *testing.T) {
|
|
h := testHub(t, Config{MaxConnections: 100})
|
|
var wg sync.WaitGroup
|
|
for i := 0; i < 50; i++ {
|
|
wg.Add(1)
|
|
go func(i int) {
|
|
defer wg.Done()
|
|
c := newConnection(fmt.Sprint(i))
|
|
if err := h.Add(c); err != nil {
|
|
t.Error(err)
|
|
return
|
|
}
|
|
for j := 0; j < 10; j++ {
|
|
_ = h.Join(c.ID(), "shared")
|
|
_, _ = h.Broadcast(ws.Message{Type: ws.Text})
|
|
h.Leave(c.ID(), "shared")
|
|
}
|
|
_ = c.Close(1000, "")
|
|
}(i)
|
|
}
|
|
wg.Wait()
|
|
_ = h.Close()
|
|
if h.Stats() != (Stats{}) {
|
|
t.Fatal("registry not released")
|
|
}
|
|
}
|
|
|
|
func TestHubRealSessionDeliveryAndDisconnect(t *testing.T) {
|
|
received := make(chan string, 1)
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
|
|
if err != nil {
|
|
return
|
|
}
|
|
defer conn.Close()
|
|
for {
|
|
_, data, err := conn.ReadMessage()
|
|
if err != nil {
|
|
return
|
|
}
|
|
received <- string(data)
|
|
}
|
|
}))
|
|
defer server.Close()
|
|
client, err := ws.NewClient(ws.ClientConfig{ProxyMode: "direct", Session: ws.SessionConfig{CloseTimeout: 100 * time.Millisecond}})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer func() {
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancel()
|
|
if err := client.Shutdown(ctx); err != nil {
|
|
t.Error(err)
|
|
}
|
|
}()
|
|
s, err := client.Dial(context.Background(), "ws"+strings.TrimPrefix(server.URL, "http"), ws.Handlers{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
h := testHub(t, Config{})
|
|
if err = h.Add(s); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err = h.Join(s.ID(), "room"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
result, err := h.SendGroup("room", ws.Message{Type: ws.Text, Data: []byte("queued")})
|
|
if err != nil || result.Queued != 1 {
|
|
t.Fatalf("send: %#v %v", result, err)
|
|
}
|
|
select {
|
|
case got := <-received:
|
|
if got != "queued" {
|
|
t.Fatal(got)
|
|
}
|
|
case <-time.After(time.Second):
|
|
t.Fatal("delivery timed out")
|
|
}
|
|
_ = s.Close(1000, "")
|
|
deadline := time.After(time.Second)
|
|
for h.Stats().Connections != 0 {
|
|
select {
|
|
case <-deadline:
|
|
t.Fatal("disconnected Session still registered")
|
|
default:
|
|
time.Sleep(time.Millisecond)
|
|
}
|
|
}
|
|
}
|
|
|
|
func BenchmarkHubBroadcast(b *testing.B) {
|
|
for _, count := range []int{1, 100, 1000} {
|
|
b.Run(fmt.Sprint(count), func(b *testing.B) {
|
|
h, _ := New(Config{})
|
|
defer h.Close()
|
|
for i := 0; i < count; i++ {
|
|
if err := h.Add(newConnection(fmt.Sprint(i))); err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
}
|
|
message := ws.Message{Type: ws.Text, Data: []byte("benchmark")}
|
|
b.ReportAllocs()
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
result, err := h.Broadcast(message)
|
|
if err != nil || result.Queued != count {
|
|
b.Fatal("broadcast failed")
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|