Skip to content

Commit

Permalink
improve ws connections
Browse files Browse the repository at this point in the history
  • Loading branch information
ayazabbas committed Nov 1, 2022
1 parent 4088c65 commit 6635a8e
Show file tree
Hide file tree
Showing 2 changed files with 3 additions and 98 deletions.
12 changes: 3 additions & 9 deletions rpc/websocket.go
Original file line number Diff line number Diff line change
Expand Up @@ -33,11 +33,10 @@ import (
)

const (
wsReadBuffer = 1024
wsWriteBuffer = 1024
wsPingInterval = 60 * time.Second
wsReadBuffer = 2048 // 1024 bytes is not enough
wsWriteBuffer = 2048 // 1024 bytes is not enough
wsPingInterval = 58 * time.Second
wsPingWriteTimeout = 5 * time.Second
wsPongTimeout = 30 * time.Second
wsMessageSizeLimit = 15 * 1024 * 1024
)

Expand Down Expand Up @@ -243,10 +242,6 @@ type websocketCodec struct {

func newWebsocketCodec(conn *websocket.Conn, host string, req http.Header) ServerCodec {
conn.SetReadLimit(wsMessageSizeLimit)
conn.SetPongHandler(func(appData string) error {
conn.SetReadDeadline(time.Time{})
return nil
})
wc := &websocketCodec{
jsonCodec: NewFuncCodec(conn, conn.WriteJSON, conn.ReadJSON).(*jsonCodec),
conn: conn,
Expand Down Expand Up @@ -306,7 +301,6 @@ func (wc *websocketCodec) pingLoop() {
wc.jsonCodec.encMu.Lock()
wc.conn.SetWriteDeadline(time.Now().Add(wsPingWriteTimeout))
wc.conn.WriteMessage(websocket.PingMessage, nil)
wc.conn.SetReadDeadline(time.Now().Add(wsPongTimeout))
wc.jsonCodec.encMu.Unlock()
timer.Reset(wsPingInterval)
}
Expand Down
89 changes: 0 additions & 89 deletions rpc/websocket_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,14 +19,10 @@ package rpc
import (
"context"
"errors"
"io"
"net"
"net/http"
"net/http/httptest"
"net/http/httputil"
"net/url"
"strings"
"sync/atomic"
"testing"
"time"

Expand Down Expand Up @@ -227,63 +223,6 @@ func TestClientWebsocketLargeMessage(t *testing.T) {
}
}

func TestClientWebsocketSevered(t *testing.T) {
t.Parallel()

var (
server = wsPingTestServer(t, nil)
ctx = context.Background()
)
defer server.Shutdown(ctx)

u, err := url.Parse("http://" + server.Addr)
if err != nil {
t.Fatal(err)
}
rproxy := httputil.NewSingleHostReverseProxy(u)
var severable *severableReadWriteCloser
rproxy.ModifyResponse = func(response *http.Response) error {
severable = &severableReadWriteCloser{ReadWriteCloser: response.Body.(io.ReadWriteCloser)}
response.Body = severable
return nil
}
frontendProxy := httptest.NewServer(rproxy)
defer frontendProxy.Close()

wsURL := "ws:" + strings.TrimPrefix(frontendProxy.URL, "http:")
client, err := DialWebsocket(ctx, wsURL, "")
if err != nil {
t.Fatalf("client dial error: %v", err)
}
defer client.Close()

resultChan := make(chan int)
sub, err := client.EthSubscribe(ctx, resultChan, "foo")
if err != nil {
t.Fatalf("client subscribe error: %v", err)
}

// sever the connection
severable.Sever()

// Wait for subscription error.
timeout := time.NewTimer(3 * wsPingInterval)
defer timeout.Stop()
for {
select {
case err := <-sub.Err():
t.Log("client subscription error:", err)
return
case result := <-resultChan:
t.Error("unexpected result:", result)
return
case <-timeout.C:
t.Error("didn't get any error within the test timeout")
return
}
}
}

// wsPingTestServer runs a WebSocket server which accepts a single subscription request.
// When a value arrives on sendPing, the server sends a ping frame, waits for a matching
// pong and finally delivers a single subscription result.
Expand Down Expand Up @@ -386,31 +325,3 @@ func wsPingTestHandler(t *testing.T, conn *websocket.Conn, shutdown, sendPing <-
}
}
}

// severableReadWriteCloser wraps an io.ReadWriteCloser and provides a Sever() method to drop writes and read empty.
type severableReadWriteCloser struct {
io.ReadWriteCloser
severed int32 // atomic
}

func (s *severableReadWriteCloser) Sever() {
atomic.StoreInt32(&s.severed, 1)
}

func (s *severableReadWriteCloser) Read(p []byte) (n int, err error) {
if atomic.LoadInt32(&s.severed) > 0 {
return 0, nil
}
return s.ReadWriteCloser.Read(p)
}

func (s *severableReadWriteCloser) Write(p []byte) (n int, err error) {
if atomic.LoadInt32(&s.severed) > 0 {
return len(p), nil
}
return s.ReadWriteCloser.Write(p)
}

func (s *severableReadWriteCloser) Close() error {
return s.ReadWriteCloser.Close()
}

0 comments on commit 6635a8e

Please sign in to comment.