proxy: warmup fixes

Author: xcolwell

proxy/main.go

@@ -102,34 +102,40 @@ func main() {
 
 	glog.Infof("Listen api (:%d), socks5 (:%d), http (:%d), https (:%d)", ListenApiPort, ListenSocksPort, ListenHttpPort, ListenHttpsPort)
 
-	newApiServer(
+	newSocks5Server(
 		ctx,
 		cancel,
 		proxyDeviceManager,
 		transportTls,
 		settings,
 	)
 
-	newSocks5Server(
+	newHttpServer(
 		ctx,
 		cancel,
 		proxyDeviceManager,
 		transportTls,
 		settings,
 	)
 
-	newHttpServer(
+	wg := newWgServer(
 		ctx,
 		cancel,
 		proxyDeviceManager,
-		transportTls,
 		settings,
 	)
 
-	wg := newWgServer(
+	warmup := func(proxyClient *model.ProxyClient) error {
+		wg.AddProxyClients(proxyClient)
+		return nil
+	}
+
+	newApiServer(
 		ctx,
 		cancel,
 		proxyDeviceManager,
+		transportTls,
+		warmup,
 		settings,
 	)
 
@@ -150,7 +156,7 @@ func main() {
 			}
 		}
 
-		wg.AddProxyClients(proxyClients)
+		wg.AddProxyClients(proxyClients...)
 	})
 	defer sub()
 
@@ -398,7 +404,7 @@ func (self *wgServer) run() {
 	}
 }
 
-func (self *wgServer) AddProxyClients(proxyClients []*model.ProxyClient) {
+func (self *wgServer) AddProxyClients(proxyClients ...*model.ProxyClient) {
 	serverConfig := model.LoadServerProxyConfig()
 
 	clients := map[netip.Addr]*proxy.WgClient{}

proxy/proxy_api.go

@@ -12,7 +12,7 @@ import (
 	"time"
 
 	"github.com/urnetwork/server"
-	// "github.com/urnetwork/server/model"
+	"github.com/urnetwork/server/model"
 	"github.com/urnetwork/server/router"
 )
 
@@ -21,6 +21,7 @@ type apiServer struct {
 	cancel             context.CancelFunc
 	proxyDeviceManager *ProxyDeviceManager
 	transportTls       *server.TransportTls
+	warmupCallback     func(*model.ProxyClient) error
 	settings           *ProxySettings
 }
 
@@ -29,13 +30,15 @@ func newApiServer(
 	cancel context.CancelFunc,
 	proxyDeviceManager *ProxyDeviceManager,
 	transportTls *server.TransportTls,
+	warmupCallback func(*model.ProxyClient) error,
 	settings *ProxySettings,
 ) *apiServer {
 	s := &apiServer{
 		ctx:                ctx,
 		cancel:             cancel,
 		proxyDeviceManager: proxyDeviceManager,
 		transportTls:       transportTls,
+		warmupCallback:     warmupCallback,
 		settings:           settings,
 	}
 
@@ -112,6 +115,21 @@ func (self *apiServer) HandleWarmup(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	if self.warmupCallback != nil {
+		proxyClient, err := model.GetProxyClient(self.ctx, proxyId)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+		if proxyClient != nil {
+			err = self.warmupCallback(proxyClient)
+			if err != nil {
+				http.Error(w, err.Error(), http.StatusInternalServerError)
+				return
+			}
+		}
+	}
+
 	timeout := time.Duration(warmupRequest.TimeoutSeconds) * time.Second
 	ready := proxyDevice.WaitForReady(r.Context(), timeout)
 

proxy/proxy_client_notification.go

@@ -17,29 +17,53 @@ import (
 type ProxyClientsFunction = func(proxyClients []*model.ProxyClient)
 
 type proxyClientNotification struct {
-	ctx context.Context
+	ctx    context.Context
+	cancel context.CancelFunc
 
 	settings *ProxySettings
 
 	proxyClientsCallbacks *connect.CallbackList[ProxyClientsFunction]
 }
 
 func newProxyClientNotification(ctx context.Context, settings *ProxySettings) *proxyClientNotification {
+	cancelCtx, cancel := context.WithCancel(ctx)
 	p := &proxyClientNotification{
-		ctx:                   ctx,
+		ctx:                   cancelCtx,
+		cancel:                cancel,
 		settings:              settings,
 		proxyClientsCallbacks: connect.NewCallbackList[ProxyClientsFunction](),
 	}
-	go server.HandleError(p.run)
+	go server.HandleError(p.run, cancel)
 	return p
 }
 
 func (self *proxyClientNotification) run() {
+	defer self.cancel()
+
 	proxyHost := server.RequireHost()
 	block := server.RequireBlock()
 
+	monitor := connect.NewMonitor()
+
+	go server.HandleError(func() {
+		defer self.cancel()
+
+		event, sub := server.Subscribe(self.ctx, model.ProxyClientChannel(proxyHost, block))
+		defer sub()
+
+		for {
+			select {
+			case <-self.ctx.Done():
+				return
+			case <-event:
+				monitor.NotifyAll()
+			}
+		}
+	})
+
 	nextChangeId := int64(0)
 	for {
+		notify := monitor.NotifyChannel()
 		proxyClients, maxChangeId, err := model.GetProxyClientsSince(
 			self.ctx,
 			proxyHost,
@@ -57,6 +81,7 @@ func (self *proxyClientNotification) run() {
 		select {
 		case <-self.ctx.Done():
 			return
+		case <-notify:
 		case <-time.After(self.settings.NotificationTimeout):
 		}
 	}