diff --git a/hub.go b/hub.go index edc6337..4a8cffa 100644 --- a/hub.go +++ b/hub.go @@ -78,3 +78,14 @@ func (h *hub) len() int { return len(h.sessions) } + +func (h *hub) all() []*Session { + h.rwmutex.RLock() + defer h.rwmutex.RUnlock() + + s := make([]*Session, 0, len(h.sessions)) + for k := range h.sessions { + s = append(s, k) + } + return s +} diff --git a/melody.go b/melody.go index e78055a..b657dde 100644 --- a/melody.go +++ b/melody.go @@ -274,6 +274,14 @@ func (m *Melody) BroadcastBinaryOthers(msg []byte, s *Session) error { }) } +// Sessions returns all sessions. An error is returned if the melody session is closed. +func (m *Melody) Sessions() ([]*Session, error) { + if m.hub.closed() { + return nil, errors.New("melody instance is closed") + } + return m.hub.all(), nil +} + // Close closes the melody instance and all connected sessions. func (m *Melody) Close() error { if m.hub.closed() { diff --git a/melody_test.go b/melody_test.go index 3136ed7..214c990 100644 --- a/melody_test.go +++ b/melody_test.go @@ -163,6 +163,56 @@ func TestLen(t *testing.T) { } } +func TestGetSessions(t *testing.T) { + rand.Seed(time.Now().UnixNano()) + + connect := int(rand.Int31n(100)) + disconnect := rand.Float32() + conns := make([]*websocket.Conn, connect) + defer func() { + for _, conn := range conns { + if conn != nil { + conn.Close() + } + } + }() + + echo := NewTestServerHandler(func(session *Session, msg []byte) {}) + server := httptest.NewServer(echo) + defer server.Close() + + disconnected := 0 + for i := 0; i < connect; i++ { + conn, err := NewDialer(server.URL) + + if err != nil { + t.Error(err) + } + + if rand.Float32() < disconnect { + conns[i] = nil + disconnected++ + conn.Close() + continue + } + + conns[i] = conn + } + + time.Sleep(time.Millisecond) + + connected := connect - disconnected + + allsess, err := echo.m.Sessions() + if err != nil { + t.Fatalf("error retrieving sessions: %v", err.Error()) + } + + if len(allsess) != connected { + t.Errorf("melody sessions %d should equal %d", len(allsess), connected) + } +} + func TestEchoBinary(t *testing.T) { echo := NewTestServer() echo.m.HandleMessageBinary(func(session *Session, msg []byte) {