Merge pull request #100 from troutowicz/per-message-deadline

Support per-message write deadlines

Author: olahol

envelope.go

@@ -1,7 +1,10 @@
 package melody
 
+import "time"
+
 type envelope struct {
-	t      int
-	msg    []byte
-	filter filterFunc
+	t         int
+	msg       []byte
+	filter    filterFunc
+	writeWait time.Duration // Optional per-message deadline (0 = use Config.WriteWait)
 }

melody.go

@@ -2,6 +2,7 @@ package melody
 
 import (
 	"net/http"
+	"time"
 
 	"github.com/gorilla/websocket"
 )
@@ -298,3 +299,73 @@ func (m *Melody) IsClosed() bool {
 func FormatCloseMessage(closeCode int, text string) []byte {
 	return websocket.FormatCloseMessage(closeCode, text)
 }
+
+// BroadcastWithDeadline broadcasts a text message with a custom write deadline.
+// If deadline is 0, uses Config.WriteWait.
+func (m *Melody) BroadcastWithDeadline(msg []byte, deadline time.Duration) error {
+	if m.hub.closed() {
+		return ErrClosed
+	}
+
+	message := envelope{
+		t:         websocket.TextMessage,
+		msg:       msg,
+		writeWait: deadline,
+	}
+	m.hub.broadcast(message)
+
+	return nil
+}
+
+// BroadcastFilterWithDeadline broadcasts a text message to filtered sessions with a custom write deadline.
+// If deadline is 0, uses Config.WriteWait.
+func (m *Melody) BroadcastFilterWithDeadline(msg []byte, deadline time.Duration, fn func(*Session) bool) error {
+	if m.hub.closed() {
+		return ErrClosed
+	}
+
+	message := envelope{
+		t:         websocket.TextMessage,
+		msg:       msg,
+		filter:    fn,
+		writeWait: deadline,
+	}
+	m.hub.broadcast(message)
+
+	return nil
+}
+
+// BroadcastBinaryWithDeadline broadcasts a binary message with a custom write deadline.
+// If deadline is 0, uses Config.WriteWait.
+func (m *Melody) BroadcastBinaryWithDeadline(msg []byte, deadline time.Duration) error {
+	if m.hub.closed() {
+		return ErrClosed
+	}
+
+	message := envelope{
+		t:         websocket.BinaryMessage,
+		msg:       msg,
+		writeWait: deadline,
+	}
+	m.hub.broadcast(message)
+
+	return nil
+}
+
+// BroadcastBinaryFilterWithDeadline broadcasts a binary message to filtered sessions with a custom write deadline.
+// If deadline is 0, uses Config.WriteWait.
+func (m *Melody) BroadcastBinaryFilterWithDeadline(msg []byte, deadline time.Duration, fn func(*Session) bool) error {
+	if m.hub.closed() {
+		return ErrClosed
+	}
+
+	message := envelope{
+		t:         websocket.BinaryMessage,
+		msg:       msg,
+		filter:    fn,
+		writeWait: deadline,
+	}
+	m.hub.broadcast(message)
+
+	return nil
+}

melody_test.go

@@ -847,3 +847,98 @@ func TestConcurrentMessageHandling(t *testing.T) {
 		}
 	})
 }
+
+func TestPerMessageDeadline(t *testing.T) {
+	done := make(chan bool)
+
+	ws := NewTestServer()
+	ws.m.Config.WriteWait = 5 * time.Second
+
+	ws.m.HandleConnect(func(s *Session) {
+		// New API with custom deadline
+		err := s.WriteWithDeadline([]byte("test"), 30*time.Second)
+		assert.Nil(t, err)
+		close(done)
+	})
+
+	server := httptest.NewServer(ws)
+	defer server.Close()
+
+	conn := MustNewDialer(server.URL)
+	defer conn.Close()
+
+	_, msg, err := conn.ReadMessage()
+	assert.Nil(t, err)
+	assert.Equal(t, "test", string(msg))
+
+	<-done
+}
+
+func TestZeroDeadlineUsesConfig(t *testing.T) {
+	done := make(chan bool)
+
+	ws := NewTestServer()
+	ws.m.Config.WriteWait = 5 * time.Second
+
+	ws.m.HandleConnect(func(s *Session) {
+		// Explicitly pass 0, should use Config.WriteWait
+		err := s.WriteWithDeadline([]byte("test"), 0)
+		assert.Nil(t, err)
+		close(done)
+	})
+
+	server := httptest.NewServer(ws)
+	defer server.Close()
+
+	conn := MustNewDialer(server.URL)
+	defer conn.Close()
+
+	_, msg, err := conn.ReadMessage()
+	assert.Nil(t, err)
+	assert.Equal(t, "test", string(msg))
+
+	<-done
+}
+
+func TestSizeBasedDeadline(t *testing.T) {
+	received := make(chan bool, 2)
+
+	ws := NewTestServer()
+	ws.m.Config.WriteWait = 10 * time.Second
+
+	server := httptest.NewServer(ws)
+	defer server.Close()
+
+	conn := MustNewDialer(server.URL)
+	defer conn.Close()
+
+	// Small message - short deadline
+	smallMsg := []byte("ping")
+	deadline := 2 * time.Second
+	err := ws.m.BroadcastWithDeadline(smallMsg, deadline)
+	assert.Nil(t, err)
+
+	// Large message - long deadline
+	largeMsg := make([]byte, 8192)
+	for i := range largeMsg {
+		largeMsg[i] = 'x'
+	}
+	longDeadline := 30 * time.Second
+	err = ws.m.BroadcastWithDeadline(largeMsg, longDeadline)
+	assert.Nil(t, err)
+
+	// Read first message
+	_, msg1, err := conn.ReadMessage()
+	assert.Nil(t, err)
+	assert.Equal(t, smallMsg, msg1)
+	received <- true
+
+	// Read second message
+	_, msg2, err := conn.ReadMessage()
+	assert.Nil(t, err)
+	assert.Equal(t, largeMsg, msg2)
+	received <- true
+
+	<-received
+	<-received
+}

session.go

@@ -39,7 +39,13 @@ func (s *Session) writeRaw(message envelope) error {
 		return ErrWriteClosed
 	}
 
-	s.conn.SetWriteDeadline(time.Now().Add(s.melody.Config.WriteWait))
+	// Use per-message deadline if specified, otherwise use global config
+	deadline := message.writeWait
+	if deadline == 0 {
+		deadline = s.melody.Config.WriteWait
+	}
+
+	s.conn.SetWriteDeadline(time.Now().Add(deadline))
 	err := s.conn.WriteMessage(message.t, message.msg)
 
 	if err != nil {
@@ -259,3 +265,35 @@ func (s *Session) RemoteAddr() net.Addr {
 func (s *Session) WebsocketConnection() *websocket.Conn {
 	return s.conn
 }
+
+// WriteWithDeadline writes a text message to the session with a custom write deadline.
+// If deadline is 0, uses Config.WriteWait.
+func (s *Session) WriteWithDeadline(msg []byte, deadline time.Duration) error {
+	if s.closed() {
+		return ErrSessionClosed
+	}
+
+	s.writeMessage(envelope{
+		t:         websocket.TextMessage,
+		msg:       msg,
+		writeWait: deadline,
+	})
+
+	return nil
+}
+
+// WriteBinaryWithDeadline writes a binary message to the session with a custom write deadline.
+// If deadline is 0, uses Config.WriteWait.
+func (s *Session) WriteBinaryWithDeadline(msg []byte, deadline time.Duration) error {
+	if s.closed() {
+		return ErrSessionClosed
+	}
+
+	s.writeMessage(envelope{
+		t:         websocket.BinaryMessage,
+		msg:       msg,
+		writeWait: deadline,
+	})
+
+	return nil
+}