Создание простого TCP чата на Go

Пример простого TCP чата на Go с объяснением логики.

Начнём с запуска сервера:

l, err := net.Listen("tcp", "localhost:9090")
if err != nil {
	return
}

defer l.Close()

Сервер слушает TCP порт 9090 на localhost. Это обычная практика — запускать сервисы, привязанные к localhost.

Теперь давайте принимаем подключения:

conn, err := l.Accept()
if err != nil {
	return
}

Пока без обработки ошибок. На этом остановимся позже.

В чате нам нужно читать сообщения, отправляемые клиентами.

Количество клиентов обычно больше одного, а чтение — это блокирующая операция. Один поток программы не позволяет нам читать из всех подключений. Исправим это, добавив горутину для каждого подключённого клиента:

l, err := net.Listen("tcp", "localhost:9090")
if err != nil {
	return
}

defer l.Close()

for {
	conn, err := l.Accept()
	if err != nil {
		return
	}

	go handleUserConnection(conn)
}

func handleUserConnection(c net.Conn) {
	defer c.Close()

	for {
		userInput, err := bufio.NewReader(c).ReadString('\n')
		if err != nil {
			return
		}
	}
}

Здесь мы возвращаемся из функции обработки при ошибке чтения, что произойдёт при отключении клиента. После возврата из функции мы также закрываем подключение. Здесь помогает defer — независимо от того, как завершается функция — defer вызывается после return.

Нам не нужно обрабатывать ошибки закрытия подключения, потому что ошибки здесь не меняют поток выполнения программы.

Теперь у нас есть работающий сервер, но нет логики, кроме чтения данных от клиентов. Добавим логику чата. Самый простой чат будет отправлять любое сообщение всем клиентам. Это можно назвать паттерном “fan out” (веерная рассылка), он проиллюстрирован ниже:

Чтобы реализовать “fan out”, нам нужно перебрать все подключения клиентов и записать в них сообщения.

Для этого нам нужно каким-то образом хранить подключения, чтобы мы могли их перебирать. Я выбрал sync.Map здесь, потому что он решает все проблемы конкурентного доступа.

Для нашей задачи хватило бы среза. Но в нашей конкурентной программе нам пришлось бы добавлять/удалять данные из этого среза, и это было бы невозможно без sync.Lock. Поскольку мы пишем простой TCP чат, давайте просто используем предопределённый тип, который решает проблемы конкурентности.

// Используем sync.Map, чтобы не иметь дело с проблемами конкурентности срезов/карт
var connMap = &sync.Map{}

Важно работать с указателем на sync.Map, а не со значением.

sync.Map — это структура, которая содержит sync.Lock. Есть известная проблема под названием “блокировки, передаваемые по значению”.

Для нашей карты нам нужны ключи. Я использую UUID здесь, потому что он гарантирует, что ключи будут разными. (Крайне низкая вероятность того, что UUID сгенерирует два одинаковых значения).

Когда клиент отключается, нам нужно удалить этого клиента из карты. Давайте передадим ID клиента в функцию обработки:

id := uuid.New().String()
connMap.Store(id, conn)

Наконец, давайте реализуем паттерн веерной рассылки сообщений:

for {
	userInput, err := bufio.NewReader(c).ReadString('\n')
	if err != nil {
		return
	}

	connMap.Range(func(key, value interface{}) bool {
		if conn, ok := value.(net.Conn); ok {
			conn.Write([]byte(userInput))
		}

		return true
	})
}

Мы можем остановить перебор карты, если вернём false из функции range.

Наконец, я добавил обработку ошибок и использование логгера zap. Следующий код — это полное решение. Полный исходный код также доступен на GitHub.

package main

import (
	"bufio"
	"net"
	"sync"

	"github.com/google/uuid"
	"go.uber.org/zap"
)

func main() {
	var loggerConfig = zap.NewProductionConfig()
	loggerConfig.Level.SetLevel(zap.DebugLevel)

	logger, err := loggerConfig.Build()
	if err != nil {
		panic(err)
	}

	l, err := net.Listen("tcp", "localhost:9090")
	if err != nil {
		return
	}

	defer l.Close()

	// Using sync.Map to not deal with concurrency slice/map issues
	var connMap = &sync.Map{}

	for {
		conn, err := l.Accept()
		if err != nil {
			logger.Error("error accepting connection", zap.Error(err))
			return
		}

		id := uuid.New().String()
		connMap.Store(id, conn)

		go handleUserConnection(id, conn, connMap, logger)
	}
}

func handleUserConnection(id string, c net.Conn, connMap *sync.Map, logger *zap.Logger) {
	defer func() {
		c.Close()
		connMap.Delete(id)
	}()

	for {
		userInput, err := bufio.NewReader(c).ReadString('\n')
		if err != nil {
			logger.Error("error reading from client", zap.Error(err))
			return
		}

		connMap.Range(func(key, value interface{}) bool {
			if conn, ok := value.(net.Conn); ok {
				if _, err := conn.Write([]byte(userInput)); err != nil {
					logger.Error("error on writing to connection", zap.Error(err))
				}
			}

			return true
		})
	}
}