refactor: worker gateway

Signed-off-by: thxCode <thxcode0824@gmail.com>

Author: thxCode

pkg/utils/httpx/middleware_logging.go

@@ -0,0 +1,101 @@
+package httpx
+
+import (
+	"bufio"
+	"errors"
+	"io"
+	"net"
+	"net/http"
+
+	klog "k8s.io/klog/v2"
+)
+
+type loggingResponseWriter struct {
+	http.ResponseWriter
+	status int
+}
+
+func (w *loggingResponseWriter) WriteHeader(code int) {
+	w.status = code
+	w.ResponseWriter.WriteHeader(code)
+}
+
+func (w *loggingResponseWriter) Write(p []byte) (int, error) {
+	if w.status == 0 {
+		w.status = http.StatusOK
+	}
+	return w.ResponseWriter.Write(p)
+}
+
+func (w *loggingResponseWriter) Flush() {
+	if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
+		flusher.Flush()
+	}
+}
+
+func (w *loggingResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	hijacker, ok := w.ResponseWriter.(http.Hijacker)
+	if !ok {
+		return nil, nil, errors.New("response writer does not support hijacking")
+	}
+	return hijacker.Hijack()
+}
+
+func (w *loggingResponseWriter) Push(target string, opts *http.PushOptions) error {
+	pusher, ok := w.ResponseWriter.(http.Pusher)
+	if !ok {
+		return http.ErrNotSupported
+	}
+	return pusher.Push(target, opts)
+}
+
+func (w *loggingResponseWriter) ReadFrom(r io.Reader) (n int64, err error) {
+	if w.status == 0 {
+		w.status = http.StatusOK
+	}
+	if readerFrom, ok := w.ResponseWriter.(io.ReaderFrom); ok {
+		return readerFrom.ReadFrom(r)
+	}
+	return io.Copy(w.ResponseWriter, r)
+}
+
+func (w *loggingResponseWriter) Unwrap() http.ResponseWriter {
+	return w.ResponseWriter
+}
+
+// AccessLog logs request method/path and response status for each request.
+func AccessLog(logger klog.Logger, queryAsValues bool) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			kvs := []any{
+				"client", r.RemoteAddr,
+				"method", r.Method,
+				"path", r.URL.Path,
+			}
+			if queryAsValues {
+				for k, v := range r.URL.Query() {
+					if k == "token" || k == "password" {
+						v = []string{"***"}
+					}
+					if len(v) == 1 {
+						kvs = append(kvs, k, v[0])
+					} else {
+						kvs = append(kvs, k, v)
+					}
+				}
+			}
+			slg := logger.WithValues(kvs...)
+			slg.Info("request")
+			rw := &loggingResponseWriter{ResponseWriter: w}
+			next.ServeHTTP(rw, r.WithContext(klog.NewContext(r.Context(), slg)))
+			if rw.status == 0 {
+				rw.status = http.StatusOK
+			}
+			if rw.status >= 400 {
+				slg.Error(nil, "response", "status", rw.status)
+			} else {
+				slg.V(2).Info("response", "status", rw.status)
+			}
+		})
+	}
+}

pkg/workergateway/manager/config.go

@@ -9,7 +9,7 @@ import (
 
 type (
 	// ConstructRestConfigFunc defines a function type for constructing rest.Config for a given cluster.
-	ConstructRestConfigFunc = func(cluster string) (*rest.Config, error)
+	ConstructRestConfigFunc = func(cluster, token string) (*rest.Config, error)
 
 	Config struct {
 		ConstructRestConfig ConstructRestConfigFunc

pkg/workergateway/manager/manager.go

@@ -68,7 +68,7 @@ type (
 	Manager interface {
 		// SubscribeWorker subscribes a worker for the given cluster.
 		// If force is true, it will unsubscribe the existing worker before subscribing a new one.
-		SubscribeWorker(ctx context.Context, cluster string, force bool) error
+		SubscribeWorker(ctx context.Context, cluster, token string, force bool) error
 		// UnsubscribeWorker unsubscribes the worker for the given cluster.
 		UnsubscribeWorker(ctx context.Context, cluster string)
 		// ListWorkers lists all subscribed workers with their status.
@@ -140,17 +140,17 @@ func New(ctx context.Context, config *Config) (Manager, error) {
 	}, nil
 }
 
-func (wm *_Manager) SubscribeWorker(ctx context.Context, cluster string, force bool) error {
+func (wm *_Manager) SubscribeWorker(ctx context.Context, cluster, token string, force bool) error {
 	logger := wm.Logger.WithValues("cluster", cluster)
 
 	if force {
 		wm.UnsubscribeWorker(ctx, cluster)
 	} else if wm.hasWorker(cluster) {
-		logger.Info("worker already exists, skip")
+		logger.V(2).Info("worker already exists, skip")
 		return nil
 	}
 
-	cfg, err := wm.ConstructRestConfig(cluster)
+	cfg, err := wm.ConstructRestConfig(cluster, token)
 	if err != nil {
 		logger.Error(err, "construct rest config")
 		return fmt.Errorf("construct rest config for cluster %q: %w", cluster, err)
@@ -172,6 +172,7 @@ func (wm *_Manager) SubscribeWorker(ctx context.Context, cluster string, force b
 			default:
 			}
 
+			logger.Info("checking api services")
 			if err := apis.WaitForServicesReady(wkCtx, cli); err != nil {
 				logger.Error(err, "wait for api services ready")
 				continue
@@ -441,7 +442,7 @@ func registerEventHandler(
 		DeleteFunc: func(obj any) { publishEvent(WorkerEventDeleted, obj) },
 	}
 	opts := cache.HandlerOptions{
-		Logger: ptr.To(klog.FromContext(ctx).WithValues("gvk", gvk)),
+		Logger: ptr.To(klog.FromContext(ctx).WithValues("gvk", gvk).V(4)),
 	}
 
 	_, _ = inf.AddEventHandlerWithOptions(handler, opts)

pkg/workergateway/manager/option.go

@@ -115,7 +115,7 @@ func (o *Options) Complete(_ context.Context) (*Config, error) {
 }
 
 func constructRestConfigByGPUStackAPI(port int, o *Options) ConstructRestConfigFunc {
-	return func(cluster string) (*rest.Config, error) {
+	return func(cluster, token string) (*rest.Config, error) {
 		scheme := o.WorkerConnGPUStackAPIScheme
 		if scheme == "" {
 			scheme = "http"
@@ -125,7 +125,8 @@ func constructRestConfigByGPUStackAPI(port int, o *Options) ConstructRestConfigF
 		}
 
 		cfg := &rest.Config{
-			Host: fmt.Sprintf("%s://localhost:%d/v2/clusters/%s/proxy", scheme, port, cluster),
+			Host:        fmt.Sprintf("%s://localhost:%d/v2/clusters/%s/proxy", scheme, port, cluster),
+			BearerToken: token,
 		}
 		cfg.UserAgent = version.GetUserAgent()
 		cfg.Timeout = o.KubeConnTimeout
@@ -152,7 +153,7 @@ func constructRestConfigByLoopback(o *Options) (ConstructRestConfigFunc, error)
 	cfg.Burst = o.KubeConnBurst
 	cfg.ContentType = o.KubeContentType
 
-	return func(_ string) (*rest.Config, error) {
+	return func(_, _ string) (*rest.Config, error) {
 		return cfg, nil
 	}, nil
 }

pkg/workergateway/service/helper.go

@@ -19,6 +19,9 @@ type ListAggregateInstanceTypes struct {
 
 func OpListAggregateInstanceTypes() *ListAggregateInstanceTypes {
 	return &ListAggregateInstanceTypes{
+		list: AggregatedInstanceTypeList{
+			Items: make([]AggregatedInstanceType, 0),
+		},
 		itemIndexer: make(map[AggregatedInstanceTypeSpec]int),
 	}
 }
@@ -359,7 +362,11 @@ type ListClusterInstanceTypes struct {
 }
 
 func OpListClusterInstanceTypes() *ListClusterInstanceTypes {
-	return &ListClusterInstanceTypes{}
+	return &ListClusterInstanceTypes{
+		list: ClusterInstanceTypeList{
+			Items: make([]ClusterInstanceType, 0),
+		},
+	}
 }
 
 func (in *ListClusterInstanceTypes) Next(cluster string, obj runtime.Object) error {
@@ -411,7 +418,11 @@ type ListClusterInstances struct {
 }
 
 func OpListClusterInstances() *ListClusterInstances {
-	return &ListClusterInstances{}
+	return &ListClusterInstances{
+		list: ClusterInstanceList{
+			Items: make([]ClusterInstance, 0),
+		},
+	}
 }
 
 func (in *ListClusterInstances) Next(cluster string, obj runtime.Object) error {
@@ -468,7 +479,11 @@ type ListClusterInstancePersistentVolumeTypes struct {
 }
 
 func OpListClusterInstancePersistentVolumeTypes() *ListClusterInstancePersistentVolumeTypes {
-	return &ListClusterInstancePersistentVolumeTypes{}
+	return &ListClusterInstancePersistentVolumeTypes{
+		list: ClusterInstancePersistentVolumeTypeList{
+			Items: make([]ClusterInstancePersistentVolumeType, 0),
+		},
+	}
 }
 
 func (in *ListClusterInstancePersistentVolumeTypes) Next(cluster string, obj runtime.Object) error {
@@ -577,7 +592,11 @@ type ListClusterInstanceImagePullSecrets struct {
 }
 
 func OpListClusterInstanceImagePullSecrets() *ListClusterInstanceImagePullSecrets {
-	return &ListClusterInstanceImagePullSecrets{}
+	return &ListClusterInstanceImagePullSecrets{
+		list: ClusterInstanceImagePullSecretList{
+			Items: make([]ClusterInstanceImagePullSecret, 0),
+		},
+	}
 }
 
 func (in *ListClusterInstanceImagePullSecrets) Next(cluster string, obj runtime.Object) error {
@@ -634,7 +653,11 @@ type ListClusterInstanceSSHPublicKeys struct {
 }
 
 func OpListClusterInstanceSSHPublicKeys() *ListClusterInstanceSSHPublicKeys {
-	return &ListClusterInstanceSSHPublicKeys{}
+	return &ListClusterInstanceSSHPublicKeys{
+		list: ClusterInstanceSSHPublicKeyList{
+			Items: make([]ClusterInstanceSSHPublicKey, 0),
+		},
+	}
 }
 
 func (in *ListClusterInstanceSSHPublicKeys) Next(cluster string, obj runtime.Object) error {