Skip to content

Commit

Permalink
Concurrent message handling
Browse files Browse the repository at this point in the history
  • Loading branch information
olahol committed Feb 12, 2024
1 parent b889a36 commit f3ccd79
Show file tree
Hide file tree
Showing 4 changed files with 104 additions and 10 deletions.
11 changes: 6 additions & 5 deletions config.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,12 @@ import "time"

// Config melody configuration struct.
type Config struct {
WriteWait time.Duration // Milliseconds until write times out.
PongWait time.Duration // Timeout for waiting on pong.
PingPeriod time.Duration // Milliseconds between pings.
MaxMessageSize int64 // Maximum size in bytes of a message.
MessageBufferSize int // The max amount of messages that can be in a sessions buffer before it starts dropping them.
WriteWait time.Duration // Duration until write times out.
PongWait time.Duration // Timeout for waiting on pong.
PingPeriod time.Duration // Duration between pings.
MaxMessageSize int64 // Maximum size in bytes of a message.
MessageBufferSize int // The max amount of messages that can be in a sessions buffer before it starts dropping them.
ConcurrentMessageHandling bool // Handle messages from sessions concurrently.
}

func newConfig() *Config {
Expand Down
5 changes: 5 additions & 0 deletions melody.go
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,11 @@ func (m *Melody) HandlePong(fn func(*Session)) {
}

// HandleMessage fires fn when a text message comes in.
// NOTE: by default Melody handles messages sequentially for each
// session. This has the effect that a message handler exceeding the
// read deadline (Config.PongWait, by default 1 minute) will time out
// the session. Concurrent message handling can be turned on by setting
// Config.ConcurrentMessageHandling to true.
func (m *Melody) HandleMessage(fn func(*Session, []byte)) {
m.messageHandler = fn
}
Expand Down
81 changes: 81 additions & 0 deletions melody_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -765,3 +765,84 @@ func TestHandleSentMessage(t *testing.T) {
conn.WriteMessage(websocket.BinaryMessage, TestMsg)
})
}

func TestConcurrentMessageHandling(t *testing.T) {
testTimeout := func(cmh bool, msgType int) bool {
base := time.Millisecond * 100
done := make(chan struct{})

handler := func(s *Session, msg []byte) {
if len(msg) == 0 {
done <- struct{}{}
return
}

time.Sleep(base * 2)
}

ws := NewTestServerHandler(func(session *Session, msg []byte) {})
if msgType == websocket.TextMessage {
ws.m.HandleMessage(handler)
} else {
ws.m.HandleMessageBinary(handler)
}

ws.m.Config.ConcurrentMessageHandling = cmh
ws.m.Config.PongWait = base

var errorSet bool
ws.m.HandleError(func(s *Session, err error) {
errorSet = true
done <- struct{}{}
})

server := httptest.NewServer(ws)
defer server.Close()

conn := MustNewDialer(server.URL)
defer conn.Close()

conn.WriteMessage(msgType, TestMsg)
conn.WriteMessage(msgType, TestMsg)

time.Sleep(base / 4)

conn.WriteMessage(msgType, nil)

<-done

return errorSet
}

t.Run("text should error", func(t *testing.T) {
errorSet := testTimeout(false, websocket.TextMessage)

if !errorSet {
t.FailNow()
}
})

t.Run("text should not error", func(t *testing.T) {
errorSet := testTimeout(true, websocket.TextMessage)

if errorSet {
t.FailNow()
}
})

t.Run("binary should error", func(t *testing.T) {
errorSet := testTimeout(false, websocket.BinaryMessage)

if !errorSet {
t.FailNow()
}
})

t.Run("binary should not error", func(t *testing.T) {
errorSet := testTimeout(true, websocket.BinaryMessage)

if errorSet {
t.FailNow()
}
})
}
17 changes: 12 additions & 5 deletions session.go
Original file line number Diff line number Diff line change
Expand Up @@ -133,13 +133,20 @@ func (s *Session) readPump() {
break
}

if t == websocket.TextMessage {
s.melody.messageHandler(s, message)
if s.melody.Config.ConcurrentMessageHandling {
go s.handleMessage(t, message)
} else {
s.handleMessage(t, message)
}
}
}

if t == websocket.BinaryMessage {
s.melody.messageHandlerBinary(s, message)
}
func (s *Session) handleMessage(t int, message []byte) {
switch t {
case websocket.TextMessage:
s.melody.messageHandler(s, message)
case websocket.BinaryMessage:
s.melody.messageHandlerBinary(s, message)
}
}

Expand Down

0 comments on commit f3ccd79

Please sign in to comment.