feat(ws): release generic WebSocket client and server v0.31.0-preview
This commit is contained in:
271
ws/hub/hub_test.go
Normal file
271
ws/hub/hub_test.go
Normal file
@@ -0,0 +1,271 @@
|
||||
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")
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user