refactor: introduce Pool interface, move agent_pool to agent module

This commit is contained in:
yusing 2025-04-13 06:11:06 +08:00
parent 12a63a66f6
commit 90214ff752
13 changed files with 132 additions and 91 deletions

16
agent/pkg/agent/agents.go Normal file
View file

@ -0,0 +1,16 @@
package agent
import (
"github.com/yusing/go-proxy/internal/utils/pool"
)
type agents struct{ pool.Pool[*AgentConfig] }
var Agents = agents{pool.New[*AgentConfig]()}
func (agents agents) Get(agentAddrOrDockerHost string) (*AgentConfig, bool) {
if !IsDockerHostAgent(agentAddrOrDockerHost) {
return agents.Base().Load(agentAddrOrDockerHost)
}
return agents.Base().Load(GetAgentAddrFromDockerHost(agentAddrOrDockerHost))
}

View file

@ -4,9 +4,9 @@ import (
"context"
"crypto/tls"
"crypto/x509"
"encoding/json"
"net"
"net/http"
"net/url"
"os"
"strings"
"time"
@ -16,7 +16,6 @@ import (
"github.com/yusing/go-proxy/internal/gperr"
"github.com/yusing/go-proxy/internal/logging"
"github.com/yusing/go-proxy/internal/net/gphttp"
"github.com/yusing/go-proxy/internal/net/types"
"github.com/yusing/go-proxy/pkg"
)
@ -49,8 +48,8 @@ const (
)
var (
AgentURL = types.MustParseURL(APIBaseURL)
HTTPProxyURL = types.MustParseURL(APIBaseURL + EndpointProxyHTTP)
AgentURL, _ = url.Parse(APIBaseURL)
HTTPProxyURL, _ = url.Parse(APIBaseURL + EndpointProxyHTTP)
HTTPProxyURLPrefixLen = len(APIEndpointBase + EndpointProxyHTTP)
)

View file

@ -22,7 +22,7 @@ func Logs(w http.ResponseWriter, r *http.Request) {
until := query.Get("to")
levels := query.Get("levels") // TODO: implement levels
dockerClient, found, err := getDockerClient(w, server)
dockerClient, found, err := getDockerClient(server)
if err != nil {
gphttp.BadRequest(w, err.Error())
return

View file

@ -8,6 +8,7 @@ import (
"github.com/coder/websocket"
"github.com/coder/websocket/wsjson"
"github.com/yusing/go-proxy/agent/pkg/agent"
config "github.com/yusing/go-proxy/internal/config/types"
"github.com/yusing/go-proxy/internal/docker"
"github.com/yusing/go-proxy/internal/gperr"
@ -44,7 +45,7 @@ func getDockerClients() (DockerClients, gperr.Error) {
dockerClients[name] = dockerClient
}
for _, agent := range cfg.ListAgents() {
for _, agent := range agent.Agents.Iter {
dockerClient, err := docker.NewClient(agent.FakeDockerHost())
if err != nil {
connErrs.Add(err)
@ -56,7 +57,7 @@ func getDockerClients() (DockerClients, gperr.Error) {
return dockerClients, connErrs.Error()
}
func getDockerClient(w http.ResponseWriter, server string) (*docker.SharedClient, bool, error) {
func getDockerClient(server string) (*docker.SharedClient, bool, error) {
cfg := config.GetInstance()
var host string
for name, h := range cfg.Value().Providers.Docker {
@ -65,7 +66,7 @@ func getDockerClient(w http.ResponseWriter, server string) (*docker.SharedClient
break
}
}
for _, agent := range cfg.ListAgents() {
for _, agent := range agent.Agents.Iter {
if agent.Name() == server {
host = agent.FakeDockerHost()
break
@ -119,6 +120,6 @@ func serveHTTP[V any, T ResultType[V]](w http.ResponseWriter, r *http.Request, g
})
} else {
result, err := getResult(r.Context(), dockerClients)
handleResult[V, T](w, err, result)
handleResult[V](w, err, result)
}
}

View file

@ -40,7 +40,7 @@ func NewAgent(w http.ResponseWriter, r *http.Request) {
return
}
hostport := fmt.Sprintf("%s:%d", host, port)
if _, ok := config.GetInstance().GetAgent(hostport); ok {
if _, ok := agent.Agents.Get(hostport); ok {
gphttp.ClientError(w, gphttp.ErrAlreadyExists("agent", hostport), http.StatusConflict)
return
}

View file

@ -3,8 +3,8 @@ package v1
import (
"net/http"
"github.com/yusing/go-proxy/agent/pkg/agent"
agentPkg "github.com/yusing/go-proxy/agent/pkg/agent"
config "github.com/yusing/go-proxy/internal/config/types"
"github.com/yusing/go-proxy/internal/gperr"
"github.com/yusing/go-proxy/internal/metrics/systeminfo"
"github.com/yusing/go-proxy/internal/net/gphttp"
@ -12,7 +12,7 @@ import (
"github.com/yusing/go-proxy/internal/net/gphttp/reverseproxy"
)
func SystemInfo(cfg config.ConfigInstance, w http.ResponseWriter, r *http.Request) {
func SystemInfo(w http.ResponseWriter, r *http.Request) {
query := r.URL.Query()
agentAddr := query.Get("agent_addr")
query.Del("agent_addr")
@ -21,7 +21,7 @@ func SystemInfo(cfg config.ConfigInstance, w http.ResponseWriter, r *http.Reques
return
}
agent, ok := cfg.GetAgent(agentAddr)
agent, ok := agent.Agents.Get(agentAddr)
if !ok {
gphttp.NotFound(w, "agent_addr")
return

View file

@ -1,66 +0,0 @@
package config
import (
"slices"
"github.com/yusing/go-proxy/agent/pkg/agent"
"github.com/yusing/go-proxy/internal/gperr"
"github.com/yusing/go-proxy/internal/route/provider"
"github.com/yusing/go-proxy/internal/utils/functional"
)
var agentPool = functional.NewMapOf[string, *agent.AgentConfig]()
func addAgent(agent *agent.AgentConfig) {
agentPool.Store(agent.Addr, agent)
}
func removeAllAgents() {
agentPool.Clear()
}
func GetAgent(addr string) (agent *agent.AgentConfig, ok bool) {
agent, ok = agentPool.Load(addr)
return
}
func (cfg *Config) GetAgent(agentAddrOrDockerHost string) (*agent.AgentConfig, bool) {
if !agent.IsDockerHostAgent(agentAddrOrDockerHost) {
return GetAgent(agentAddrOrDockerHost)
}
return GetAgent(agent.GetAgentAddrFromDockerHost(agentAddrOrDockerHost))
}
func (cfg *Config) VerifyNewAgent(host string, ca agent.PEMPair, client agent.PEMPair) (int, gperr.Error) {
if slices.ContainsFunc(cfg.value.Providers.Agents, func(a *agent.AgentConfig) bool {
return a.Addr == host
}) {
return 0, gperr.New("agent already exists")
}
var agentCfg agent.AgentConfig
agentCfg.Addr = host
err := agentCfg.InitWithCerts(cfg.task.Context(), ca.Cert, client.Cert, client.Key)
if err != nil {
return 0, gperr.Wrap(err, "failed to start agent")
}
addAgent(&agentCfg)
provider := provider.NewAgentProvider(&agentCfg)
if err := cfg.errIfExists(provider); err != nil {
return 0, err
}
err = provider.LoadRoutes()
if err != nil {
return 0, gperr.Wrap(err, "failed to load routes")
}
return provider.NumRoutes(), nil
}
func (cfg *Config) ListAgents() []*agent.AgentConfig {
agents := make([]*agent.AgentConfig, 0, agentPool.Size())
agentPool.RangeAll(func(key string, value *agent.AgentConfig) {
agents = append(agents, value)
})
return agents
}

View file

@ -277,8 +277,8 @@ func (cfg *Config) errIfExists(p *proxy.Provider) gperr.Error {
func (cfg *Config) initAgents(agentCfgs []*agent.AgentConfig) gperr.Error {
var wg sync.WaitGroup
var errs gperr.Builder
errs := gperr.NewBuilderWithConcurrency()
wg.Add(len(agentCfgs))
for _, agentCfg := range agentCfgs {
go func(agentCfg *agent.AgentConfig) {
@ -286,7 +286,7 @@ func (cfg *Config) initAgents(agentCfgs []*agent.AgentConfig) gperr.Error {
if err := agentCfg.Init(cfg.task.Context()); err != nil {
errs.Add(err.Subject(agentCfg.String()))
} else {
addAgent(agentCfg)
agent.Agents.Add(agentCfg)
}
}(agentCfg)
}
@ -298,7 +298,7 @@ func (cfg *Config) loadRouteProviders(providers *config.Providers) gperr.Error {
errs := gperr.NewBuilder("route provider errors")
results := gperr.NewBuilder("loaded route providers")
removeAllAgents()
agent.Agents.Clear()
n := len(providers.Agents) + len(providers.Docker) + len(providers.Files)
if n == 0 {
@ -309,12 +309,12 @@ func (cfg *Config) loadRouteProviders(providers *config.Providers) gperr.Error {
errs.Add(cfg.initAgents(providers.Agents))
for _, agent := range providers.Agents {
if !agent.IsInitialized() { // failed to initialize
for _, a := range providers.Agents {
if !a.IsInitialized() { // failed to initialize
continue
}
addAgent(agent)
routeProviders = append(routeProviders, proxy.NewAgentProvider(agent))
agent.Agents.Add(a)
routeProviders = append(routeProviders, proxy.NewAgentProvider(a))
}
for _, filename := range providers.Files {
routeProviders = append(routeProviders, proxy.NewFileProvider(filename))
@ -338,6 +338,8 @@ func (cfg *Config) loadRouteProviders(providers *config.Providers) gperr.Error {
lenLongestName = len(k)
}
})
errs.EnableConcurrency()
results.EnableConcurrency()
cfg.providers.RangeAllParallel(func(_ string, p *proxy.Provider) {
if err := p.LoadRoutes(); err != nil {
errs.Add(err.Subject(p.String()))

View file

@ -1,6 +1,10 @@
package config
import (
"slices"
"github.com/yusing/go-proxy/agent/pkg/agent"
"github.com/yusing/go-proxy/internal/gperr"
"github.com/yusing/go-proxy/internal/route"
"github.com/yusing/go-proxy/internal/route/provider"
)
@ -51,3 +55,29 @@ func (cfg *Config) Statistics() map[string]any {
"providers": providerStats,
}
}
func (cfg *Config) VerifyNewAgent(host string, ca agent.PEMPair, client agent.PEMPair) (int, gperr.Error) {
if slices.ContainsFunc(cfg.value.Providers.Agents, func(a *agent.AgentConfig) bool {
return a.Addr == host
}) {
return 0, gperr.New("agent already exists")
}
var agentCfg agent.AgentConfig
agentCfg.Addr = host
err := agentCfg.InitWithCerts(cfg.task.Context(), ca.Cert, client.Cert, client.Key)
if err != nil {
return 0, gperr.Wrap(err, "failed to start agent")
}
agent.Agents.Add(&agentCfg)
provider := provider.NewAgentProvider(&agentCfg)
if err := cfg.errIfExists(provider); err != nil {
return 0, err
}
err = provider.LoadRoutes()
if err != nil {
return 0, gperr.Wrap(err, "failed to load routes")
}
return provider.NumRoutes(), nil
}

View file

@ -45,9 +45,7 @@ type (
Statistics() map[string]any
RouteProviderList() []string
Context() context.Context
GetAgent(agentAddrOrDockerHost string) (*agent.AgentConfig, bool)
VerifyNewAgent(host string, ca agent.PEMPair, client agent.PEMPair) (int, gperr.Error)
ListAgents() []*agent.AgentConfig
AutoCertProvider() *autocert.Provider
}
)

View file

@ -14,7 +14,6 @@ import (
"github.com/docker/docker/client"
"github.com/yusing/go-proxy/agent/pkg/agent"
"github.com/yusing/go-proxy/internal/common"
config "github.com/yusing/go-proxy/internal/config/types"
"github.com/yusing/go-proxy/internal/logging"
"github.com/yusing/go-proxy/internal/task"
)
@ -134,7 +133,7 @@ func NewClient(host string) (*SharedClient, error) {
var dial func(ctx context.Context) (net.Conn, error)
if agent.IsDockerHostAgent(host) {
cfg, ok := config.GetInstance().GetAgent(host)
cfg, ok := agent.Agents.Get(host)
if !ok {
panic(fmt.Errorf("agent %q not found", host))
}

View file

@ -83,7 +83,7 @@ func FromDocker(c *container.Summary, dockerHost string) (res *Container) {
if agent.IsDockerHostAgent(dockerHost) {
var ok bool
res.Agent, ok = config.GetInstance().GetAgent(dockerHost)
res.Agent, ok = agent.Agents.Get(dockerHost)
if !ok {
logging.Error().Msgf("agent %q not found", dockerHost)
}

View file

@ -0,0 +1,62 @@
package pool
import (
"sort"
"github.com/yusing/go-proxy/internal/utils"
"github.com/yusing/go-proxy/internal/utils/functional"
)
type (
Pool[T Object] struct {
m functional.Map[string, T]
}
Object interface {
Key() string
Name() string
utils.MapMarshaller
}
)
func New[T Object]() Pool[T] {
return Pool[T]{functional.NewMapOf[string, T]()}
}
func (p Pool[T]) Add(obj T) {
p.m.Store(obj.Key(), obj)
}
func (p Pool[T]) Get(key string) (T, bool) {
return p.m.Load(key)
}
func (p Pool[T]) Size() int {
return p.m.Size()
}
func (p Pool[T]) Clear() {
p.m.Clear()
}
func (p Pool[T]) Base() functional.Map[string, T] {
return p.m
}
func (p Pool[T]) Slice() []T {
slice := make([]T, 0, p.m.Size())
for _, v := range p.m.Range {
slice = append(slice, v)
}
sort.Slice(slice, func(i, j int) bool {
return slice[i].Name() < slice[j].Name()
})
return slice
}
func (p Pool[T]) Iter(fn func(k string, v T) bool) {
p.m.Range(fn)
}
func (p Pool[T]) IterAll(fn func(k string, v T)) {
p.m.RangeAll(fn)
}