Replace gorilla ws

.github/workflows/test.yml

@@ -13,7 +13,7 @@ jobs:
  test:
  strategy:
  matrix:
- go-version: ['1.21']
+ go-version: ["1.21"]
  permissions:
  contents: read
  pull-requests: read
@@ -28,7 +28,7 @@ jobs:
  - name: download surrealdb
  run: curl --proto '=https' --tlsv1.2 -sSf https://install.surrealdb.com | sh -s -- --nightly
  - name: start surrealdb
- run: surreal start memory -A --auth --user root --pass root &
+ run: surreal start memory -A --user root --pass root &
  - name: test
  run: go test -v -cover ./...
  env:

db_test.go

@@ -9,7 +9,6 @@ import (
  "os"
  "sync"
  "testing"
- "time"
 
  "github.com/stretchr/testify/assert"
  "github.com/surrealdb/surrealdb.go/pkg/logger/slog"
@@ -18,7 +17,7 @@ import (
  "github.com/stretchr/testify/require"
  "github.com/stretchr/testify/suite"
  "github.com/surrealdb/surrealdb.go"
- "github.com/surrealdb/surrealdb.go/pkg/conn/gorilla"
+ "github.com/surrealdb/surrealdb.go/pkg/conn/nhooyr"
  "github.com/surrealdb/surrealdb.go/pkg/constants"
 
  "github.com/surrealdb/surrealdb.go/pkg/conn"
@@ -65,17 +64,18 @@ func TestSurrealDBSuite(t *testing.T) {
  SurrealDBSuite := new(SurrealDBTestSuite)
  SurrealDBSuite.connImplementations = make(map[string]conn.Connection)
 
+ // // Nhooyr
  // Without options
- buff := bytes.NewBufferString("")
- logData := createLogger(t, buff)
- SurrealDBSuite.connImplementations["gorilla"] = gorilla.Create().Logger(logData)
- SurrealDBSuite.logBuffer = buff
+ nbuff := bytes.NewBufferString("")
+ nlogData := createLogger(t, nbuff)
+ SurrealDBSuite.connImplementations["nhooyr"] = nhooyr.Create().Logger(nlogData)
+ SurrealDBSuite.logBuffer = nbuff
 
  // With options
- buffOpt := bytes.NewBufferString("")
- logDataOpt := createLogger(t, buff)
- SurrealDBSuite.connImplementations["gorilla_opt"] = gorilla.Create().SetTimeOut(time.Minute).SetCompression(true).Logger(logDataOpt)
- SurrealDBSuite.logBuffer = buffOpt
+ nbuffOpt := bytes.NewBufferString("")
+ nlogDataOpt := createLogger(t, nbuffOpt)
+ SurrealDBSuite.connImplementations["nhooyr_opt"] = nhooyr.Create().Logger(nlogDataOpt)
+ SurrealDBSuite.logBuffer = nbuffOpt
 
  RunWsMap(t, SurrealDBSuite)
 }
@@ -781,7 +781,7 @@ func (s *SurrealDBTestSuite) TestConcurrentOperations() {
 }
 
 func (s *SurrealDBTestSuite) TestConnectionBreak() {
- ws := gorilla.Create()
+ ws := nhooyr.Create()
  var url string
  if currentURL == "" {
  url = defaultURL
@@ -791,10 +791,11 @@ func (s *SurrealDBTestSuite) TestConnectionBreak() {
 
  db := s.openConnection(url, ws)
  // Close the connection hard from ws
- ws.Conn.Close()
+ err := ws.Conn.CloseNow()
+ s.Require().NoError(err)
 
  // Needs to be return error when the connection is closed or broken
- _, err := db.Select("users")
+ _, err = db.Select("users")
  s.Require().Error(err)
 }
 

go.mod

@@ -3,8 +3,10 @@ module github.com/surrealdb/surrealdb.go
 go 1.20
 
 require (
+ github.com/coder/websocket v1.8.12
  github.com/gorilla/websocket v1.5.0
  github.com/stretchr/testify v1.8.4
+ nhooyr.io/websocket v1.8.17
 )
 
 require (

go.sum

@@ -1,3 +1,5 @@
+github.com/coder/websocket v1.8.12 h1:5bUXkEPPIbewrnkU8LTCLVaxi4N4J8ahufH2vlo4NAo=
+github.com/coder/websocket v1.8.12/go.mod h1:LNVeNrXQZfe5qhS9ALED3uA+l5pPqvwXg3CKoDBB2gs=
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -23,3 +25,7 @@ gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntN
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+nhooyr.io/websocket v1.8.10 h1:mv4p+MnGrLDcPlBoWsvPP7XCzTYMXP9F9eIGoKbgx7Q=
+nhooyr.io/websocket v1.8.10/go.mod h1:rN9OFWIUwuxg4fR5tELlYC04bXYowCP9GX47ivo2l+c=
+nhooyr.io/websocket v1.8.17 h1:KEVeLJkUywCKVsnLIDlD/5gtayKp8VoCkksHCGGfT9Y=
+nhooyr.io/websocket v1.8.17/go.mod h1:rN9OFWIUwuxg4fR5tELlYC04bXYowCP9GX47ivo2l+c=

pkg/conn/gorilla/gorilla.go → pkg/conn/nhooyr/nhooyr.go

@@ -1,19 +1,20 @@
-package gorilla
+package nhooyr
 
 import (
+ "context"
  "encoding/json"
  "errors"
  "fmt"
- "io"
  "net"
+ "net/http"
  "reflect"
  "strconv"
  "sync"
  "time"
 
+ coderws "github.com/coder/websocket"
  "github.com/surrealdb/surrealdb.go/pkg/model"
 
- gorilla "github.com/gorilla/websocket"
  "github.com/surrealdb/surrealdb.go/internal/rpc"
  "github.com/surrealdb/surrealdb.go/pkg/conn"
  "github.com/surrealdb/surrealdb.go/pkg/logger"
@@ -32,7 +33,7 @@ const (
 type Option func(ws *WebSocket) error
 
 type WebSocket struct {
- Conn *gorilla.Conn
+ Conn *coderws.Conn
  connLock sync.Mutex
  Timeout time.Duration
  Option []Option
@@ -44,29 +45,33 @@ type WebSocket struct {
  notificationChannels map[string]chan model.Notification
  notificationChannelsLock sync.RWMutex
 
- closeChan chan int
- closeError error
+ close chan int
 }
 
 func Create() *WebSocket {
  return &WebSocket{
  Conn: nil,
- closeChan: make(chan int),
+ close:  make(chan int),
  responseChannels: make(map[string]chan rpc.RPCResponse),
  notificationChannels: make(map[string]chan model.Notification),
  Timeout: DefaultTimeout * time.Second,
  }
 }
 
 func (ws *WebSocket) Connect(url string) (conn.Connection, error) {
- dialer := gorilla.DefaultDialer
- dialer.EnableCompression = true
+ ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+ defer cancel()
 
- connection, _, err := dialer.Dial(url, nil)
+ connection, resp, err := coderws.Dial(ctx, url, nil)
  if err != nil {
  return nil, err
  }
 
+ if resp.StatusCode != http.StatusSwitchingProtocols {
+ resp.Body.Close()
+ return nil, fmt.Errorf("unexpected status code: %d", resp.StatusCode)
+ }
+
  ws.Conn = connection
 
  for _, option := range ws.Option {
@@ -75,7 +80,7 @@ func (ws *WebSocket) Connect(url string) (conn.Connection, error) {
  }
  }
 
- go ws.initialize()
+ ws.initialize()
  return ws, nil
 }
 
@@ -98,24 +103,12 @@ func (ws *WebSocket) RawLogger(logData logger.Logger) *WebSocket {
  return ws
 }
 
-func (ws *WebSocket) SetCompression(compress bool) *WebSocket {
- ws.Option = append(ws.Option, func(ws *WebSocket) error {
- ws.Conn.EnableWriteCompression(compress)
- return nil
- })
- return ws
-}
-
 func (ws *WebSocket) Close() error {
  ws.connLock.Lock()
  defer ws.connLock.Unlock()
- close(ws.closeChan)
- err := ws.Conn.WriteMessage(gorilla.CloseMessage, gorilla.FormatCloseMessage(CloseMessageCode, ""))
- if err != nil {
- return err
- }
+ close(ws.close)
 
- return ws.Conn.Close()
+ return ws.Conn.Close(coderws.StatusNormalClosure, "")
 }
 
 func (ws *WebSocket) LiveNotifications(liveQueryID string) (chan model.Notification, error) {
@@ -181,12 +174,6 @@ func (ws *WebSocket) getLiveChannel(id string) (chan model.Notification, bool) {
 }
 
 func (ws *WebSocket) Send(method string, params []interface{}) (interface{}, error) {
- select {
- case <-ws.closeChan:
- return nil, ws.closeError
- default:
- }
-
  id := rand.String(RequestIDLength)
  request := &rpc.RPCRequest{
  ID: id,
@@ -224,7 +211,7 @@ func (ws *WebSocket) Send(method string, params []interface{}) (interface{}, err
 }
 
 func (ws *WebSocket) read(v interface{}) error {
- _, data, err := ws.Conn.ReadMessage()
+ _, data, err := ws.Conn.Read(context.Background())
  if err != nil {
  return err
  }
@@ -239,42 +226,29 @@ func (ws *WebSocket) write(v interface{}) error {
 
  ws.connLock.Lock()
  defer ws.connLock.Unlock()
- return ws.Conn.WriteMessage(gorilla.TextMessage, data)
+ return ws.Conn.Write(context.Background(), coderws.MessageText, data)
 }
 
 func (ws *WebSocket) initialize() {
- for {
- select {
- case <-ws.closeChan:
- return
- default:
- var res rpc.RPCResponse
- err := ws.read(&res)
- if err != nil {
- shouldExit := ws.handleError(err)
- if shouldExit {
- return
+ go func() {
+ for {
+ select {
+ case <-ws.close:
+ return
+ default:
+ var res rpc.RPCResponse
+ err := ws.read(&res)
+ if err != nil {
+ if errors.Is(err, net.ErrClosed) {
+ break
+ }
+ ws.logger.Error(err.Error())
+ continue
  }
- continue
+ go ws.handleResponse(res)
  }
- go ws.handleResponse(res)
  }
- }
-}
-
-func (ws *WebSocket) handleError(err error) bool {
- if errors.Is(err, net.ErrClosed) {
- ws.closeError = net.ErrClosed
- return true
- }
- if gorilla.IsUnexpectedCloseError(err) {
- ws.closeError = io.ErrClosedPipe
- <-ws.closeChan
- return true
- }
-
- ws.logger.Error(err.Error())
- return false
+ }()
 }
 
 func (ws *WebSocket) handleResponse(res rpc.RPCResponse) {