GoDoxy/internal/route/stream.go

package route

import (
	"context"
	"errors"

	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"
	"github.com/yusing/go-proxy/internal/gperr"
	"github.com/yusing/go-proxy/internal/idlewatcher"
	net "github.com/yusing/go-proxy/internal/net/types"
	"github.com/yusing/go-proxy/internal/route/routes"
	"github.com/yusing/go-proxy/internal/task"
	"github.com/yusing/go-proxy/internal/watcher/health"
	"github.com/yusing/go-proxy/internal/watcher/health/monitor"
)

// TODO: support stream load balance.
type StreamRoute struct {
	*Route

	net.Stream `json:"-"`

	HealthMon health.HealthMonitor `json:"health"`

	task *task.Task

	l zerolog.Logger
}

func NewStreamRoute(base *Route) (routes.Route, gperr.Error) {
	// TODO: support non-coherent scheme
	return &StreamRoute{
		Route: base,
		l: log.With().
			Str("type", string(base.Scheme)).
			Str("name", base.Name()).
			Logger(),
	}, nil
}

// Start implements task.TaskStarter.
func (r *StreamRoute) Start(parent task.Parent) gperr.Error {
	r.task = parent.Subtask("stream."+r.Name(), !r.ShouldExclude())
	r.Stream = NewStream(r)

	switch {
	case r.UseIdleWatcher():
		waker, err := idlewatcher.NewWatcher(parent, r, r.IdlewatcherConfig())
		if err != nil {
			r.task.Finish(err)
			return gperr.Wrap(err, "idlewatcher error")
		}
		r.Stream = waker
		r.HealthMon = waker
	case r.UseHealthCheck():
		r.HealthMon = monitor.NewMonitor(r)
	}

	if !r.ShouldExclude() {
		if err := r.Setup(); err != nil {
			r.task.Finish(err)
			return gperr.Wrap(err)
		}
		r.l.Info().Int("port", r.Port.Listening).Msg("listening")
	}

	if r.HealthMon != nil {
		if err := r.HealthMon.Start(r.task); err != nil {
			gperr.LogWarn("health monitor error", err, &r.l)
		}
	}

	if r.ShouldExclude() {
		return nil
	}

	if err := checkExists(r); err != nil {
		return err
	}

	go r.acceptConnections()

	routes.Stream.Add(r)
	r.task.OnCancel("remove_route_from_stream", func() {
		routes.Stream.Del(r)
	})
	return nil
}

// Task implements task.TaskStarter.
func (r *StreamRoute) Task() *task.Task {
	return r.task
}

// Finish implements task.TaskFinisher.
func (r *StreamRoute) Finish(reason any) {
	r.task.Finish(reason)
}

func (r *StreamRoute) HealthMonitor() health.HealthMonitor {
	return r.HealthMon
}

func (r *StreamRoute) acceptConnections() {
	defer r.task.Finish("listener closed")

	go func() {
		<-r.task.Context().Done()
		r.Close()
	}()

	for {
		select {
		case <-r.task.Context().Done():
			return
		default:
			conn, err := r.Accept()
			if err != nil {
				select {
				case <-r.task.Context().Done():
				default:
					gperr.LogError("accept connection error", err, &r.l)
				}
				r.task.Finish(err)
				return
			}
			if conn == nil {
				panic("connection is nil")
			}
			go func() {
				err := r.Handle(conn)
				if err != nil && !errors.Is(err, context.Canceled) {
					gperr.LogError("handle connection error", err, &r.l)
				}
			}()
		}
	}
}