В этом небольшом туториале, мы чуть подробнее разберем использование Gorilla WebSocket для написания своего websocket сервера, на примере чуть более функциональном, чем базовый пример и более легком для понимания, чем пример чата.

Что будет уметь наш сервер?

  1. Отправлять новые сообщения от клиентов в callback

  2. Хранить активные соединения и закрывать/удалять не активные

  3. Рассылать сообщения по активным соединениям

Для начала поднимем обычный http сервер при помощи net/http, для того чтобы мы могли отлавливать запросы на соединение:

package main

import (
	"fmt"
	"html"
	"net/http"
)

func main() {
	http.HandleFunc("/", echo)
	http.ListenAndServe(":8080", nil)
}

func echo(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "Hello, %q", html.EscapeString(r.URL.Path))
}

Теперь научим его "апгрейдить" соединение:

import (
	"github.com/gorilla/websocket"
	"net/http"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true // Пропускаем любой запрос
	},
}

func echo(w http.ResponseWriter, r *http.Request) {
	connection, _ := upgrader.Upgrade(w, r, nil)
	defer connection.Close() // Закрываем соединение
}

Теперь у нас есть соединение с клиентом, которые мы сразу же закрываем. Мы можем циклично читать сообщения, которые нам шлет клиент и отправлять их обратно:

func echo(w http.ResponseWriter, r *http.Request) {
	connection, _ := upgrader.Upgrade(w, r, nil)

	for {
		_, message, _ := connection.ReadMessage()

    connection.WriteMessage(websocket.TextMessage, message)
		go messageHandler(message)
	}
}

func messageHandler(message []byte)  {
  fmt.Println(string(message))
}

Научим наш сервер закрывать соединение:

func echo(w http.ResponseWriter, r *http.Request) {
  connection, _ := upgrader.Upgrade(w, r, nil)
  defer connection.Close()

	for {
		mt, message, err := connection.ReadMessage()

		if err != nil || mt == websocket.CloseMessage {
			break // Выходим из цикла, если клиент пытается закрыть соединение или связь с клиентом прервана
		}

		connection.WriteMessage(websocket.TextMessage, message)

		go messageHandler(message)
	}
}

Чтобы иметь возможность рассылать сообщения по разным соединениям, нам нужно где то их хранить, в нашем случае подойдет простейший map:

var clients map[*websocket.Conn]bool

func echo(w http.ResponseWriter, r *http.Request) {
	connection, _ := upgrader.Upgrade(w, r, nil)
  defer connection.Close()

	clients[connection] = true // Сохраняем соединение, используя его как ключ
  defer delete(clients, connection) // Удаляем соединение
  
	for {
		mt, message, err := connection.ReadMessage()

		if err != nil || mt == websocket.CloseMessage {
			break // Выходим из цикла, если клиент пытается закрыть соединение или связь прервана
		}

		// Теперь мы рассылаем сообщения всем клиентам
		go writeMessage(message)

		go messageHandler(message)
	}
}

func writeMessage(message []byte) {
	for conn := range clients {
		conn.WriteMessage(websocket.TextMessage, message)
	}
}

Теперь мы можем запаковать наш сервер в структуру, чтобы иметь возможность рассылать и принимать сообщения из вне:

package ws

import (
	"github.com/gorilla/websocket"
	"net/http"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true // Пропускаем любой запрос
	},
}

type Server struct {
	clients       map[*websocket.Conn]bool
	handleMessage func(message []byte) // хандлер новых сообщений
}

func StartServer(handleMessage func(message []byte)) *Server {
	server := Server{
		make(map[*websocket.Conn]bool),
		handleMessage,
	}

	http.HandleFunc("/", server.echo)
	go http.ListenAndServe(":8080", nil) // Уводим http сервер в горутину 

	return &server
}

func (server *Server) echo(w http.ResponseWriter, r *http.Request) {
	connection, _ := upgrader.Upgrade(w, r, nil)
  defer connection.Close()

	server.clients[connection] = true // Сохраняем соединение, используя его как ключ
  defer delete(server.clients, connection) // Удаляем соединение

	for {
		mt, message, err := connection.ReadMessage()

		if err != nil || mt == websocket.CloseMessage {
			break // Выходим из цикла, если клиент пытается закрыть соединение или связь прервана
		}

		go server.handleMessage(message)
	}
}

func (server *Server) WriteMessage(message []byte) {
	for conn := range server.clients {
		conn.WriteMessage(websocket.TextMessage, message)
	}
}
package main

import (
	"fmt"
	"simple-webcoket/ws"
)

func main() {
	server := ws.StartServer(messageHandler)

	for {
		server.WriteMessage([]byte("Hello"))
	}
}

func messageHandler(message []byte) {
	fmt.Println(string(message))
}

Теперь у нас есть реализация простейшего webscoket сервера, который способен принимать и рассылать сообщения по активным соединениям.

Комментарии (7)


  1. falconandy
    09.01.2022 17:49

    Навскидку пара замечаний по коду:

    1. Работа с clients не потокобезопасна.
    2. Закрытие соединения и удаление из clients — использовать defer:

    connection, _ := upgrader.Upgrade(w, r, nil)
    defer connection.Close()
    
    server.clients[connection] = true // Сохраняем соединение, используя его как ключ
    defer delete(server.clients, connection)
    


    1. DavidNadejdin Автор
      09.01.2022 20:43

      спасибо, учту. map был использован, так как суть туториала показать базовую работу с библиотекой, чуть лучше и проще чем в официальной документации, не больше. А насчет defer, это такое требование к код стайлу или все таки вкусовщина?


    1. DavidNadejdin Автор
      09.01.2022 20:50

      Просто задаваясь вопросом, как отслеживать все что происходит по окончанию функции, очевидно что посмотреть на все defer. Но как таким образом контролировать, в каком порядке производятся те или иные действия


      1. falconandy
        10.01.2022 09:22
        +1

        1. defer срабатывают в обратном порядке
        2. Если вы добавите return где-то в середине функции или случится panic, то код в defer всё равно отработает
        3. В другом языке можно было бы использовать try/finally — по сравнению с defer код «очистки» находится далеко от кода «создания», а также добавляется дополнительный (дополнительные) уровень вложенности


        1. DavidNadejdin Автор
          10.01.2022 12:38

          Спасибо, учел это, скорректировал код в статье


  1. z0ic
    09.01.2022 21:57

    Мне gobwas/ws больше понравился, хотя может я просто ещё не напоролся на подводные камни.


    1. falconandy
      10.01.2022 09:22

      Я еще использовал nhooyr/websocket