go-redis/tcp/echo.go

// Author: simon (ynwdlxm@163.com)
// Date: 2025/5/26 10:52
// Desc:

package tcp

import (
	"bufio"
	"context"
	"io"
	"net"
	"sync"
	"time"

	"github.com/runningwater/go-redis/lib/logger"
	"github.com/runningwater/go-redis/lib/sync/atomic"
	"github.com/runningwater/go-redis/lib/sync/wait"
)

// EchoClient 客户端信息
type EchoClient struct {
	Conn    net.Conn
	Waiting wait.Wait
}

func (e *EchoClient) Close() error {
	e.Waiting.WaitWithTimeout(10 * time.Second)
	_ = e.Conn.Close()
	return nil
}

type Echo struct {
	activeConn sync.Map
	closing    atomic.Boolean
}

func NewEcho() *Echo {
	return &Echo{
		activeConn: sync.Map{},
		closing:    0,
	}
}

func (e *Echo) Handle(ctx context.Context, conn net.Conn) {
	if e.closing.Get() {
		_ = conn.Close()
	}
	// 记录客户端信息
	client := &EchoClient{
		Conn: conn,
	}
	e.activeConn.Store(client, struct{}{})

	reader := bufio.NewReader(conn)
	for {
		msg, err := reader.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				logger.Info("Connection close")
				e.activeConn.Delete(client)
			} else {
				logger.Warn(err)
			}
			return
		}

		client.Waiting.Add(1)
		b := []byte(msg)
		_, _ = conn.Write(b)
		client.Waiting.Done()
	}
}

func (e *Echo) Close() error {
	logger.Info("handler shutting down")
	e.closing.Set(true)

	// 剔除所有连接
	e.activeConn.Range(func(key, value any) bool {
		client := key.(*EchoClient)
		_ = client.Conn.Close()
		return true
	})

	return nil
}