ws

Documentation

Overview

Example

package main

import (
	"fmt"
	"log"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"

	"github.com/gorilla/websocket"

	"github.com/determined-ai/determined/master/pkg/ws"
)

func main() {
	// Start a Websocket server that converts ints to strings.
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		upgrader := websocket.Upgrader{}
		c, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			log.Println(err)
			return
		}

		// Send and receive messages from the client.
		s, err := ws.Wrap[int, string]("int2str", c)
		if err != nil {
			log.Println(err)
			return
		}
		for {
			num, ok := <-s.Inbox
			if !ok {
				log.Println(s.Error())
				return
			}

			select {
			case s.Outbox <- strconv.Itoa(num):
			case <-s.Done:
				log.Println(s.Error())
				return
			}
		}
	}))
	defer ts.Close()

	// Connect a Websocket to the server with using `github.com/gorilla/websocket`.
	//nolint: bodyclose
	c, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(ts.URL, "http"), nil)
	if err != nil {
		log.Println(err)
		return
	}

	// Use Wrap to wrap the underlying connection.
	s, err := ws.Wrap[string, int]("client", c)
	if err != nil {
		log.Println(err)
		return
	}
	defer func() {
		// Gracefully close our connection, using Close. If it fails, it will forcibly clean up.
		if err := s.Close(); err != nil {
			log.Println(err)
			return
		}
	}()

	// Send and receive messages to the server.
	select {
	case s.Outbox <- 42:
	case <-s.Done:
		log.Println(s.Error())
		return
	}

	str, ok := <-s.Inbox
	if !ok {
		log.Println(s.Error())
		return
	}
	fmt.Println(str)
}

Output:

42

Index

• func UpgradeConnection(resp http.ResponseWriter, req *http.Request) (*websocket.Conn, error)
• func UpgradeEchoConnection(e echo.Context) (*websocket.Conn, error)
• type WebSocket
  • func Wrap[TIn, TOut any](name string, conn *websocket.Conn) (*WebSocket[TIn, TOut], error)
  • func (s *WebSocket[TIn, TOut]) Close() error
  • func (s *WebSocket[TIn, TOut]) Error() error
  • func (s *WebSocket[TIn, TOut]) Name() string

Examples

• Package

Functions

func UpgradeConnection

func UpgradeConnection(resp http.ResponseWriter, req *http.Request) (*websocket.Conn, error)

UpgradeConnection is a helper function for upgrading stdlib HTTP requests to WebSockets using a zero-value Gorilla Upgrader.

func UpgradeEchoConnection

func UpgradeEchoConnection(e echo.Context) (*websocket.Conn, error)

UpgradeEchoConnection is a helper function for upgrading Echo requests to WebSockets using a zero-value Gorilla Upgrader.

Types

type WebSocket

type WebSocket[TIn, TOut any] struct {

	// Inbox is a channel for incoming messages.
	Inbox <-chan TIn
	// Outbox is a channel for outgoing messages.
	Outbox chan<- TOut
	// Done notifies when the Websocket is closed. A read on Done blocks until this condition.
	Done <-chan struct{}
	// contains filtered or unexported fields
}

WebSocket is a facade that wraps a Gorilla websocket and provides a higher-level, type-safe, and thread-safe API by specializing for JSON encoding/decoding and using channels for read/write. The Close method must be called or resources will be leaked.

func Wrap

func Wrap[TIn, TOut any](name string, conn *websocket.Conn) (*WebSocket[TIn, TOut], error)

Wrap the given, underlying *websocket.Conn and returns a higher level, thread-safe wrapper.

func (*WebSocket[TIn, TOut]) Close

func (s *WebSocket[TIn, TOut]) Close() error

Close closes the WebSocket by performing the close handshake and closing the underlying connection, rendering it unusable.

func (*WebSocket[TIn, TOut]) Error

func (s *WebSocket[TIn, TOut]) Error() error

Error returns an error if the Websocket has encountered one. Errors from closing are excluded.

func (*WebSocket[TIn, TOut]) Name

func (s *WebSocket[TIn, TOut]) Name() string

Name grants access to the WebSocket's name.