diff --git a/api/v1alpha1/amaltheasession_children.go b/api/v1alpha1/amaltheasession_children.go index 5e8d9a8f..4dc959df 100644 --- a/api/v1alpha1/amaltheasession_children.go +++ b/api/v1alpha1/amaltheasession_children.go @@ -282,6 +282,10 @@ func (cr *AmaltheaSession) Service() v1.Service { }, } if cr.Spec.SessionLocation == Remote { + // NOTE: In order to connect through the reverse tunnel, + // we need to have the tunnel established so we publish + // the service without the pod being ready + svc.Spec.PublishNotReadyAddresses = true svc.Spec.Ports = append(svc.Spec.Ports, v1.ServicePort{ Protocol: v1.ProtocolTCP, Name: tunnelServiceName, @@ -991,6 +995,18 @@ func (cr *AmaltheaSession) sessionContainerRemote(volumeMounts []v1.VolumeMount) Name: "RSC_SERVER_PORT", Value: fmt.Sprintf("%d", RemoteSessionControllerPort), }, + v1.EnvVar{ + Name: "RSC_SESSION_PORT", + Value: fmt.Sprintf("%d", cr.Spec.Session.Port), + }, + v1.EnvVar{ + Name: "RSC_SESSION_URL_PATH", + Value: cr.Spec.Session.URLPath, + }, + v1.EnvVar{ + Name: "RSC_READINESS_PROBE_TYPE", + Value: string(cr.Spec.Session.ReadinessProbe.Type), + }, v1.EnvVar{ Name: "RSC_WSTUNNEL_SECRET", ValueFrom: ptr.To(v1.EnvVarSource{ diff --git a/internal/remote/config/config.go b/internal/remote/config/config.go index e14eb722..c14af31d 100644 --- a/internal/remote/config/config.go +++ b/internal/remote/config/config.go @@ -18,6 +18,8 @@ limitations under the License. package config import ( + "fmt" + amaltheadevv1alpha1 "github.com/SwissDataScienceCenter/amalthea/api/v1alpha1" "github.com/spf13/cobra" "github.com/spf13/viper" @@ -35,8 +37,11 @@ const ( ) const ( - serverPortFlag = "server-port" - fakeStartFlag = "fake-start" + serverPortFlag = "server-port" + fakeStartFlag = "fake-start" + sessionPortFlag = "session-port" + sessionURLPathFlag = "session-url-path" + readinessProbeTypeFlag = "readiness-probe-type" ) type RemoteSessionControllerConfig struct { @@ -53,6 +58,15 @@ type RemoteSessionControllerConfig struct { // FakeStart if true, do not start the remote session and print debug information FakeStart bool + + // SessionPort is the port where the remote session is expected to be serving + SessionPort int32 + + // SessionURLPath is the URL path for the HTTP readiness probe + SessionURLPath string + + // ReadinessProbeType is "none", "tcp", or "http" + ReadinessProbeType string } func SetFlags(cmd *cobra.Command) error { @@ -72,6 +86,30 @@ func SetFlags(cmd *cobra.Command) error { return err } + cmd.Flags().Int32(sessionPortFlag, 0, "port the remote session is expected to be serving on") + if err := viper.BindPFlag(sessionPortFlag, cmd.Flags().Lookup(sessionPortFlag)); err != nil { + return err + } + if err := viper.BindEnv(sessionPortFlag, configUtils.AsEnvVarFlag(sessionPortFlag)); err != nil { + return err + } + + cmd.Flags().String(sessionURLPathFlag, "/", "URL path for the HTTP readiness probe") + if err := viper.BindPFlag(sessionURLPathFlag, cmd.Flags().Lookup(sessionURLPathFlag)); err != nil { + return err + } + if err := viper.BindEnv(sessionURLPathFlag, configUtils.AsEnvVarFlag(sessionURLPathFlag)); err != nil { + return err + } + + cmd.Flags().String(readinessProbeTypeFlag, "none", "readiness probe type: none, tcp, or http") + if err := viper.BindPFlag(readinessProbeTypeFlag, cmd.Flags().Lookup(readinessProbeTypeFlag)); err != nil { + return err + } + if err := viper.BindEnv(readinessProbeTypeFlag, configUtils.AsEnvVarFlag(readinessProbeTypeFlag)); err != nil { + return err + } + // Set up shared flags if err := configUtils.SetFlags(cmd); err != nil { return err @@ -97,11 +135,20 @@ func GetConfig() (cfg RemoteSessionControllerConfig, err error) { cfg.ServerPort = viper.GetInt32(serverPortFlag) cfg.FakeStart = viper.GetBool(fakeStartFlag) + cfg.SessionPort = viper.GetInt32(sessionPortFlag) + cfg.SessionURLPath = viper.GetString(sessionURLPathFlag) + cfg.ReadinessProbeType = viper.GetString(readinessProbeTypeFlag) return cfg, nil } func (cfg *RemoteSessionControllerConfig) Validate() error { + if cfg.ReadinessProbeType != string(amaltheadevv1alpha1.None) && + cfg.ReadinessProbeType != string(amaltheadevv1alpha1.TCP) && + cfg.ReadinessProbeType != string(amaltheadevv1alpha1.HTTP) { + return fmt.Errorf("invalid readiness probe type: %s", cfg.ReadinessProbeType) + } + // FireCREST has priority over Runai cfg.RemoteKind = RemoteKindFirecrest firecrestConfigErr := cfg.Firecrest.Validate() diff --git a/internal/remote/config/config_test.go b/internal/remote/config/config_test.go index 77e541f3..f69e68eb 100644 --- a/internal/remote/config/config_test.go +++ b/internal/remote/config/config_test.go @@ -19,6 +19,7 @@ package config import ( "testing" + amaltheadevv1alpha1 "github.com/SwissDataScienceCenter/amalthea/api/v1alpha1" "github.com/spf13/cobra" "github.com/spf13/viper" "github.com/stretchr/testify/assert" @@ -86,3 +87,112 @@ func TestConfig(t *testing.T) { }) } } + +func TestConfigNewFlags(t *testing.T) { + tests := []struct { + name string + args []string + env map[string]string + wantProbeType string + wantSessionPort int32 + wantURLPath string + wantErr bool + }{ + { + name: "default with empty readiness probe", + args: []string{ + "--firecrest-api-url=https://firecrest.example.com", + "--firecrest-system-name=test-system", + "--auth-kind=client_credentials", + "--auth-token-uri=https://auth.example.com/token", + "--auth-firecrest-client-id=my-client", + "--auth-firecrest-client-secret=my-secret", + }, + wantProbeType: string(amaltheadevv1alpha1.None), + wantSessionPort: 0, + wantURLPath: "/", + }, + { + name: "custom flags via CLI", + args: []string{ + "--readiness-probe-type=tcp", + "--session-port=8888", + "--session-url-path=/lab", + "--runai-base-url=https://runai.example.com", + "--runai-project=my-project", + "--auth-kind=client_credentials", + "--auth-token-uri=https://runai-auth.example.com/token", + "--auth-runai-client-id=runai-client", + "--auth-runai-client-secret=runai-secret", + }, + wantProbeType: string(amaltheadevv1alpha1.TCP), + wantSessionPort: 8888, + wantURLPath: "/lab", + }, + { + name: "custom flags via environment variables", + args: []string{ + "--firecrest-api-url=https://firecrest.example.com", + "--firecrest-system-name=test-system", + "--auth-kind=client_credentials", + "--auth-token-uri=https://auth.example.com/token", + "--auth-firecrest-client-id=my-client", + "--auth-firecrest-client-secret=my-secret", + }, + env: map[string]string{ + "RSC_READINESS_PROBE_TYPE": "http", + "RSC_SESSION_PORT": "9999", + "RSC_SESSION_URL_PATH": "/jupyter", + }, + wantProbeType: string(amaltheadevv1alpha1.HTTP), + wantSessionPort: 9999, + wantURLPath: "/jupyter", + }, + { + name: "invalid readiness probe type errors", + args: []string{ + "--readiness-probe-type=invalid", + "--firecrest-api-url=https://firecrest.example.com", + "--firecrest-system-name=test-system", + "--auth-kind=client_credentials", + "--auth-token-uri=https://auth.example.com/token", + "--auth-firecrest-client-id=my-client", + "--auth-firecrest-client-secret=my-secret", + }, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + viper.Reset() + for k, v := range tt.env { + t.Setenv(k, v) + } + + cmd := &cobra.Command{ + Use: "test", + Run: func(cmd *cobra.Command, args []string) {}, + } + err := SetFlags(cmd) + require.NoError(t, err) + + cmd.SetArgs(tt.args) + err = cmd.Execute() + require.NoError(t, err) + + cfg, err := GetConfig() + require.NoError(t, err) + + err = cfg.Validate() + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + assert.Equal(t, tt.wantProbeType, cfg.ReadinessProbeType) + assert.Equal(t, tt.wantSessionPort, cfg.SessionPort) + assert.Equal(t, tt.wantURLPath, cfg.SessionURLPath) + }) + } +} diff --git a/internal/remote/server/server.go b/internal/remote/server/server.go index 363cf237..227f3d8c 100644 --- a/internal/remote/server/server.go +++ b/internal/remote/server/server.go @@ -20,11 +20,13 @@ import ( "context" "fmt" "log/slog" + "net" "net/http" "os" "os/signal" "time" + amaltheadevv1alpha1 "github.com/SwissDataScienceCenter/amalthea/api/v1alpha1" "github.com/SwissDataScienceCenter/amalthea/internal/common" "github.com/SwissDataScienceCenter/amalthea/internal/remote/config" "github.com/SwissDataScienceCenter/amalthea/internal/remote/controller" @@ -56,7 +58,7 @@ func Start() { os.Exit(1) } - server := newServer(controller) + server := newServer(controller, cfg) address := fmt.Sprintf(":%d", cfg.ServerPort) @@ -95,7 +97,7 @@ func Start() { var logLevel *slog.LevelVar = new(slog.LevelVar) var jsonLogger *slog.Logger = slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: logLevel})) -func newServer(controller controller.RemoteSessionController) (server *echo.Echo) { +func newServer(controller controller.RemoteSessionController, cfg config.RemoteSessionControllerConfig) (server *echo.Echo) { e := echo.New() e.HideBanner = true @@ -114,7 +116,49 @@ func newServer(controller controller.RemoteSessionController) (server *echo.Echo // Readiness endpoint e.GET("/ready", func(c echo.Context) error { - return c.NoContent(http.StatusOK) + switch cfg.ReadinessProbeType { + case string(amaltheadevv1alpha1.TCP): + dialer := net.Dialer{} + dialCtx, dialCtxCancel := context.WithTimeout(c.Request().Context(), 5*time.Second) + defer dialCtxCancel() + conn, err := dialer.DialContext(dialCtx, "tcp", fmt.Sprintf("127.0.0.1:%d", cfg.SessionPort)) + + if err != nil { + return c.NoContent(http.StatusServiceUnavailable) + } + if err := conn.Close(); err != nil { + slog.Error("failed to close readiness probe connection", "error", err) + } + return c.NoContent(http.StatusOK) + case string(amaltheadevv1alpha1.HTTP): + client := &http.Client{ + Timeout: 5 * time.Second, + CheckRedirect: func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + }, + } + url := fmt.Sprintf("http://127.0.0.1:%d%s", cfg.SessionPort, cfg.SessionURLPath) + req, err := http.NewRequestWithContext(c.Request().Context(), "GET", url, nil) + if err != nil { + return c.NoContent(http.StatusServiceUnavailable) + } + resp, err := client.Do(req) + if err != nil { + return c.NoContent(http.StatusServiceUnavailable) + } + defer func() { + if err := resp.Body.Close(); err != nil { + slog.Error("failed to close readiness probe response body", "error", err) + } + }() + if resp.StatusCode >= 200 && resp.StatusCode < 400 { + return c.NoContent(http.StatusOK) + } + return c.NoContent(http.StatusServiceUnavailable) + default: + // Case None or unset: preserve old behavior for backward compatibility + return c.NoContent(http.StatusOK) + } }) // Status endpoint diff --git a/internal/remote/server/server_test.go b/internal/remote/server/server_test.go new file mode 100644 index 00000000..3a5b6026 --- /dev/null +++ b/internal/remote/server/server_test.go @@ -0,0 +1,200 @@ +package server + +import ( + "context" + "fmt" + "net" + "net/http" + "net/http/httptest" + "testing" + + amaltheadevv1alpha1 "github.com/SwissDataScienceCenter/amalthea/api/v1alpha1" + "github.com/SwissDataScienceCenter/amalthea/internal/remote/config" + "github.com/SwissDataScienceCenter/amalthea/internal/remote/models" + "github.com/stretchr/testify/assert" +) + +type mockController struct{} + +func (m *mockController) Status(ctx context.Context) (models.RemoteSessionState, error) { + return models.Running, nil +} +func (m *mockController) Start(ctx context.Context) error { return nil } +func (m *mockController) Stop(ctx context.Context) error { return nil } + +func getFreePortOrDie() int32 { + var a *net.TCPAddr + var err error + if a, err = net.ResolveTCPAddr("tcp", "localhost:0"); err == nil { + var l *net.TCPListener + if l, err = net.ListenTCP("tcp", a); err == nil { + defer func() { + err := l.Close() + if err != nil { + panic(err) + } + }() + return int32(l.Addr().(*net.TCPAddr).Port) + } + panic(err) + } + panic(err) +} + +func TestReadyEndpoint(t *testing.T) { + baseCfg := config.RemoteSessionControllerConfig{ + ServerPort: 65532, + SessionPort: 0, + SessionURLPath: "/", + ReadinessProbeType: "none", + } + + tests := []struct { + name string + makeCfg func(port int32) config.RemoteSessionControllerConfig + setupBackend func(port int32) (cleanup func()) + wantStatus int + }{ + { + name: "none probe returns 200", + makeCfg: func(int32) config.RemoteSessionControllerConfig { + c := baseCfg + c.ReadinessProbeType = string(amaltheadevv1alpha1.None) + return c + }, + wantStatus: http.StatusOK, + }, + { + name: "interactive tcp open port returns 200", + makeCfg: func(port int32) config.RemoteSessionControllerConfig { + c := baseCfg + c.ReadinessProbeType = string(amaltheadevv1alpha1.TCP) + c.SessionPort = port + return c + }, + setupBackend: func(port int32) func() { + ln, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port)) + if err != nil { + panic(err) + } + return func() { _ = ln.Close() } + }, + wantStatus: http.StatusOK, + }, + { + name: "interactive tcp closed port returns 503", + makeCfg: func(port int32) config.RemoteSessionControllerConfig { + c := baseCfg + c.ReadinessProbeType = string(amaltheadevv1alpha1.TCP) + c.SessionPort = port + return c + }, + wantStatus: http.StatusServiceUnavailable, + }, + { + name: "interactive http serving 200 returns 200", + makeCfg: func(port int32) config.RemoteSessionControllerConfig { + c := baseCfg + c.ReadinessProbeType = string(amaltheadevv1alpha1.HTTP) + c.SessionPort = port + c.SessionURLPath = "/lab" + return c + }, + setupBackend: func(port int32) func() { + ln, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port)) + if err != nil { + panic(err) + } + srv := &http.Server{ + Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/lab" { + w.WriteHeader(http.StatusOK) + return + } + w.WriteHeader(http.StatusNotFound) + }), + } + go func() { _ = srv.Serve(ln) }() + return func() { _ = srv.Close() } + }, + wantStatus: http.StatusOK, + }, + { + name: "interactive http redirect 302 returns 200", + makeCfg: func(port int32) config.RemoteSessionControllerConfig { + c := baseCfg + c.ReadinessProbeType = string(amaltheadevv1alpha1.HTTP) + c.SessionPort = port + c.SessionURLPath = "/" + return c + }, + setupBackend: func(port int32) func() { + ln, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port)) + if err != nil { + panic(err) + } + srv := &http.Server{ + Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, "/other", http.StatusFound) + }), + } + go func() { _ = srv.Serve(ln) }() + return func() { _ = srv.Close() } + }, + wantStatus: http.StatusOK, + }, + { + name: "interactive http serving 500 returns 503", + makeCfg: func(port int32) config.RemoteSessionControllerConfig { + c := baseCfg + c.ReadinessProbeType = string(amaltheadevv1alpha1.HTTP) + c.SessionPort = port + c.SessionURLPath = "/" + return c + }, + setupBackend: func(port int32) func() { + ln, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port)) + if err != nil { + panic(err) + } + srv := &http.Server{ + Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + }), + } + go func() { _ = srv.Serve(ln) }() + return func() { _ = srv.Close() } + }, + wantStatus: http.StatusServiceUnavailable, + }, + { + name: "interactive http connection refused returns 503", + makeCfg: func(port int32) config.RemoteSessionControllerConfig { + c := baseCfg + c.ReadinessProbeType = string(amaltheadevv1alpha1.HTTP) + c.SessionPort = port + c.SessionURLPath = "/" + return c + }, + wantStatus: http.StatusServiceUnavailable, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + port := getFreePortOrDie() + cfg := tt.makeCfg(port) + + if tt.setupBackend != nil { + cleanup := tt.setupBackend(port) + defer cleanup() + } + + e := newServer(&mockController{}, cfg) + req := httptest.NewRequest(http.MethodGet, "/ready", nil) + rec := httptest.NewRecorder() + e.ServeHTTP(rec, req) + assert.Equal(t, tt.wantStatus, rec.Code) + }) + } +}