allow to customize readiness probe path and configure additional comp…

…onent ports (ytsaurus#251) * allow to customize readiness probe path * linting * fix * stage * stage * refactor * simplify
leo-astorsky · Aug 13, 2024 · b194ae2 · b194ae2
1 parent 4433fd9
commit b194ae2
Show file tree

Hide file tree

Showing 18 changed files with 165 additions and 15 deletions.
diff --git a/pkg/components/controller_agent.go b/pkg/components/controller_agent.go
@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"go.ytsaurus.tech/library/go/ptr"
+	corev1 "k8s.io/api/core/v1"
 
 	ytv1 "github.com/ytsaurus/yt-k8s-operator/api/v1"
 	"github.com/ytsaurus/yt-k8s-operator/pkg/apiproxy"
@@ -41,6 +42,11 @@ func NewControllerAgent(cfgen *ytconfig.Generator, ytsaurus *apiproxy.Ytsaurus,
 		"ca",
 		"controller-agents",
 		func() ([]byte, error) { return cfgen.GetControllerAgentConfig(resource.Spec.ControllerAgents) },
+		WithContainerPorts(corev1.ContainerPort{
+			Name:          consts.YTRPCPortName,
+			ContainerPort: consts.ControllerAgentRPCPort,
+			Protocol:      corev1.ProtocolTCP,
+		}),
 	)
 
 	return &ControllerAgent{

diff --git a/pkg/components/data_node.go b/pkg/components/data_node.go
@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"go.ytsaurus.tech/library/go/ptr"
+	corev1 "k8s.io/api/core/v1"
 
 	ytv1 "github.com/ytsaurus/yt-k8s-operator/api/v1"
 	"github.com/ytsaurus/yt-k8s-operator/pkg/apiproxy"
@@ -48,6 +49,11 @@ func NewDataNode(
 		func() ([]byte, error) {
 			return cfgen.GetDataNodeConfig(spec)
 		},
+		WithContainerPorts(corev1.ContainerPort{
+			Name:          consts.YTRPCPortName,
+			ContainerPort: consts.DataNodeRPCPort,
+			Protocol:      corev1.ProtocolTCP,
+		}),
 	)
 
 	return &DataNode{

diff --git a/pkg/components/discovery.go b/pkg/components/discovery.go
@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"go.ytsaurus.tech/library/go/ptr"
+	corev1 "k8s.io/api/core/v1"
 
 	ytv1 "github.com/ytsaurus/yt-k8s-operator/api/v1"
 	"github.com/ytsaurus/yt-k8s-operator/pkg/apiproxy"
@@ -42,6 +43,11 @@ func NewDiscovery(cfgen *ytconfig.Generator, ytsaurus *apiproxy.Ytsaurus) *Disco
 		func() ([]byte, error) {
 			return cfgen.GetDiscoveryConfig(&resource.Spec.Discovery)
 		},
+		WithContainerPorts(corev1.ContainerPort{
+			Name:          consts.YTRPCPortName,
+			ContainerPort: consts.DiscoveryRPCPort,
+			Protocol:      corev1.ProtocolTCP,
+		}),
 	)
 
 	return &Discovery{

diff --git a/pkg/components/exec_node.go b/pkg/components/exec_node.go
@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"go.ytsaurus.tech/library/go/ptr"
+	corev1 "k8s.io/api/core/v1"
 
 	ytv1 "github.com/ytsaurus/yt-k8s-operator/api/v1"
 	"github.com/ytsaurus/yt-k8s-operator/pkg/apiproxy"
@@ -47,6 +48,11 @@ func NewExecNode(
 		func() ([]byte, error) {
 			return cfgen.GetExecNodeConfig(spec)
 		},
+		WithContainerPorts(corev1.ContainerPort{
+			Name:          consts.YTRPCPortName,
+			ContainerPort: consts.ExecNodeRPCPort,
+			Protocol:      corev1.ProtocolTCP,
+		}),
 	)
 
 	var sidecarConfig *ConfigHelper

diff --git a/pkg/components/exec_node_remote.go b/pkg/components/exec_node_remote.go
@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"go.ytsaurus.tech/library/go/ptr"
+	corev1 "k8s.io/api/core/v1"
 
 	ytv1 "github.com/ytsaurus/yt-k8s-operator/api/v1"
 	"github.com/ytsaurus/yt-k8s-operator/pkg/apiproxy"
@@ -47,6 +48,11 @@ func NewRemoteExecNodes(
 		func() ([]byte, error) {
 			return cfgen.GetExecNodeConfig(spec)
 		},
+		WithContainerPorts(corev1.ContainerPort{
+			Name:          consts.YTRPCPortName,
+			ContainerPort: consts.ExecNodeRPCPort,
+			Protocol:      corev1.ProtocolTCP,
+		}),
 	)
 	return &RemoteExecNode{
 		baseComponent: baseComponent{labeller: &l},

diff --git a/pkg/components/httpproxy.go b/pkg/components/httpproxy.go
@@ -4,7 +4,6 @@ import (
 	"context"
 
 	"go.ytsaurus.tech/library/go/ptr"
-
 	corev1 "k8s.io/api/core/v1"
 
 	ytv1 "github.com/ytsaurus/yt-k8s-operator/api/v1"
@@ -56,6 +55,25 @@ func NewHTTPProxy(
 		func() ([]byte, error) {
 			return cfgen.GetHTTPProxyConfig(spec)
 		},
+		WithContainerPorts(
+			corev1.ContainerPort{
+				Name:          consts.YTRPCPortName,
+				ContainerPort: consts.HTTPProxyRPCPort,
+				Protocol:      corev1.ProtocolTCP,
+			},
+			corev1.ContainerPort{
+				Name:          "http",
+				ContainerPort: consts.HTTPProxyHTTPPort,
+				Protocol:      corev1.ProtocolTCP,
+			},
+			corev1.ContainerPort{
+				Name:          "https",
+				ContainerPort: consts.HTTPProxyHTTPSPort,
+				Protocol:      corev1.ProtocolTCP,
+			},
+		),
+		WithCustomReadinessProbeEndpointPort(consts.HTTPProxyHTTPPort),
+		WithCustomReadinessProbeEndpointPath("/ping"),
 	)
 
 	var httpsSecret *resources.TLSSecret

diff --git a/pkg/components/master.go b/pkg/components/master.go
@@ -55,6 +55,11 @@ func NewMaster(cfgen *ytconfig.Generator, ytsaurus *apiproxy.Ytsaurus) *Master {
 		cfgen.GetMastersStatefulSetName(),
 		cfgen.GetMastersServiceName(),
 		func() ([]byte, error) { return cfgen.GetMasterConfig(&resource.Spec.PrimaryMasters) },
+		WithContainerPorts(corev1.ContainerPort{
+			Name:          consts.YTRPCPortName,
+			ContainerPort: consts.MasterRPCPort,
+			Protocol:      corev1.ProtocolTCP,
+		}),
 	)
 
 	initJob := NewInitJob(

diff --git a/pkg/components/master_caches.go b/pkg/components/master_caches.go
@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"go.ytsaurus.tech/library/go/ptr"
+	corev1 "k8s.io/api/core/v1"
 
 	ytv1 "github.com/ytsaurus/yt-k8s-operator/api/v1"
 	"github.com/ytsaurus/yt-k8s-operator/pkg/apiproxy"
@@ -41,6 +42,11 @@ func NewMasterCache(cfgen *ytconfig.Generator, ytsaurus *apiproxy.Ytsaurus) *Mas
 		cfgen.GetMasterCachesStatefulSetName(),
 		cfgen.GetMasterCachesServiceName(),
 		func() ([]byte, error) { return cfgen.GetMasterCachesConfig(resource.Spec.MasterCaches) },
+		WithContainerPorts(corev1.ContainerPort{
+			Name:          consts.YTRPCPortName,
+			ContainerPort: consts.MasterCachesRPCPort,
+			Protocol:      corev1.ProtocolTCP,
+		}),
 	)
 
 	return &MasterCache{

diff --git a/pkg/components/query_tracker.go b/pkg/components/query_tracker.go
@@ -59,6 +59,11 @@ func NewQueryTracker(
 		cfgen.GetQueryTrackerStatefulSetName(),
 		cfgen.GetQueryTrackerServiceName(),
 		func() ([]byte, error) { return cfgen.GetQueryTrackerConfig(resource.Spec.QueryTrackers) },
+		WithContainerPorts(corev1.ContainerPort{
+			Name:          consts.YTRPCPortName,
+			ContainerPort: consts.QueryTrackerRPCPort,
+			Protocol:      corev1.ProtocolTCP,
+		}),
 	)
 
 	image := ytsaurus.GetResource().Spec.CoreImage

diff --git a/pkg/components/queue_agent.go b/pkg/components/queue_agent.go
@@ -61,6 +61,11 @@ func NewQueueAgent(
 		cfgen.GetQueueAgentStatefulSetName(),
 		cfgen.GetQueueAgentServiceName(),
 		func() ([]byte, error) { return cfgen.GetQueueAgentConfig(resource.Spec.QueueAgents) },
+		WithContainerPorts(corev1.ContainerPort{
+			Name:          consts.YTRPCPortName,
+			ContainerPort: consts.QueueAgentRPCPort,
+			Protocol:      corev1.ProtocolTCP,
+		}),
 	)
 
 	image := ytsaurus.GetResource().Spec.CoreImage

diff --git a/pkg/components/rpcproxy.go b/pkg/components/rpcproxy.go
@@ -53,6 +53,11 @@ func NewRPCProxy(
 		func() ([]byte, error) {
 			return cfgen.GetRPCProxyConfig(spec)
 		},
+		WithContainerPorts(corev1.ContainerPort{
+			Name:          consts.YTRPCPortName,
+			ContainerPort: consts.RPCProxyRPCPort,
+			Protocol:      corev1.ProtocolTCP,
+		}),
 	)
 
 	var balancingService *resources.RPCService = nil

diff --git a/pkg/components/scheduler.go b/pkg/components/scheduler.go
@@ -59,6 +59,11 @@ func NewScheduler(
 		func() ([]byte, error) {
 			return cfgen.GetSchedulerConfig(resource.Spec.Schedulers)
 		},
+		WithContainerPorts(corev1.ContainerPort{
+			Name:          consts.YTRPCPortName,
+			ContainerPort: consts.SchedulerRPCPort,
+			Protocol:      corev1.ProtocolTCP,
+		}),
 	)
 
 	return &Scheduler{

diff --git a/pkg/components/server.go b/pkg/components/server.go
@@ -20,7 +20,6 @@ import (
 )
 
 const (
-	serverAPIPortName      = "http"
 	readinessProbeHTTPPath = "/orchid/service"
 )
 
@@ -55,6 +54,11 @@ type serverImpl struct {
 	configHelper      *ConfigHelper
 
 	builtStatefulSet *appsv1.StatefulSet
+
+	componentContainerPorts []corev1.ContainerPort
+
+	readinessProbePort     intstr.IntOrString
+	readinessProbeHTTPPath string
 }
 
 func newServer(
@@ -63,6 +67,7 @@ func newServer(
 	instanceSpec *ytv1.InstanceSpec,
 	binaryPath, configFileName, statefulSetName, serviceName string,
 	generator ytconfig.YsonGeneratorFunc,
+	options ...Option,
 ) server {
 	proxy := ytsaurus.APIProxy()
 	commonSpec := ytsaurus.GetCommonSpec()
@@ -73,6 +78,7 @@ func newServer(
 		instanceSpec,
 		binaryPath, configFileName, statefulSetName, serviceName,
 		generator,
+		options...,
 	)
 }
 
@@ -83,6 +89,7 @@ func newServerConfigured(
 	instanceSpec *ytv1.InstanceSpec,
 	binaryPath, configFileName, statefulSetName, serviceName string,
 	generator ytconfig.YsonGeneratorFunc,
+	optFuncs ...Option,
 ) server {
 	image := commonSpec.CoreImage
 	if instanceSpec.Image != nil {
@@ -107,6 +114,22 @@ func newServerConfigured(
 			consts.BusSecretMountPoint)
 	}
 
+	opts := &options{
+		containerPorts: []corev1.ContainerPort{
+			{
+				Name:          consts.YTMonitoringContainerPortName,
+				ContainerPort: *instanceSpec.MonitoringPort,
+				Protocol:      corev1.ProtocolTCP,
+			},
+		},
+		readinessProbeEndpointPort: intstr.FromString(consts.YTMonitoringContainerPortName),
+		readinessProbeEndpointPath: readinessProbeHTTPPath,
+	}
+
+	for _, fn := range optFuncs {
+		fn(opts)
+	}
+
 	return &serverImpl{
 		labeller:     l,
 		image:        image,
@@ -143,6 +166,11 @@ func newServerConfigured(
 					Fmt: ytconfig.ConfigFormatYson,
 				},
 			}),
+
+		componentContainerPorts: opts.containerPorts,
+
+		readinessProbePort:     opts.readinessProbeEndpointPort,
+		readinessProbeHTTPPath: opts.readinessProbeEndpointPath,
 	}
 }
 
@@ -272,19 +300,13 @@ func (s *serverImpl) rebuildStatefulSet() *appsv1.StatefulSet {
 				Name:         consts.YTServerContainerName,
 				Command:      command,
 				VolumeMounts: volumeMounts,
-				Ports: []corev1.ContainerPort{
-					{
-						Name:          serverAPIPortName,
-						ContainerPort: *s.instanceSpec.MonitoringPort,
-						Protocol:      corev1.ProtocolTCP,
-					},
-				},
-				Resources: s.instanceSpec.Resources,
+				Ports:        s.componentContainerPorts,
+				Resources:    s.instanceSpec.Resources,
 				ReadinessProbe: &corev1.Probe{
 					ProbeHandler: corev1.ProbeHandler{
 						HTTPGet: &corev1.HTTPGetAction{
-							Port: intstr.FromString(serverAPIPortName),
-							Path: readinessProbeHTTPPath,
+							Port: s.readinessProbePort,
+							Path: s.readinessProbeHTTPPath,
 						},
 					},
 					InitialDelaySeconds: readinessProbeParams.InitialDelaySeconds,

diff --git a/pkg/components/serveroptions.go b/pkg/components/serveroptions.go
@@ -0,0 +1,33 @@
+package components
+
+import (
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/util/intstr"
+)
+
+type options struct {
+	containerPorts []corev1.ContainerPort
+
+	readinessProbeEndpointPort intstr.IntOrString
+	readinessProbeEndpointPath string
+}
+
+type Option func(opts *options)
+
+func WithCustomReadinessProbeEndpointPort(port int32) Option {
+	return func(opts *options) {
+		opts.readinessProbeEndpointPort = intstr.FromInt32(port)
+	}
+}
+
+func WithCustomReadinessProbeEndpointPath(path string) Option {
+	return func(opts *options) {
+		opts.readinessProbeEndpointPath = path
+	}
+}
+
+func WithContainerPorts(ports ...corev1.ContainerPort) Option {
+	return func(opts *options) {
+		opts.containerPorts = append(opts.containerPorts, ports...)
+	}
+}
diff --git a/pkg/components/tablet_node.go b/pkg/components/tablet_node.go
@@ -7,6 +7,7 @@ import (
 	"go.ytsaurus.tech/library/go/ptr"
 	"go.ytsaurus.tech/yt/go/ypath"
 	"go.ytsaurus.tech/yt/go/yt"
+	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 
@@ -62,6 +63,11 @@ func NewTabletNode(
 		func() ([]byte, error) {
 			return cfgen.GetTabletNodeConfig(spec)
 		},
+		WithContainerPorts(corev1.ContainerPort{
+			Name:          consts.YTRPCPortName,
+			ContainerPort: consts.TabletNodeRPCPort,
+			Protocol:      corev1.ProtocolTCP,
+		}),
 	)
 
 	return &TabletNode{

diff --git a/pkg/components/yql_agent.go b/pkg/components/yql_agent.go
@@ -51,6 +51,11 @@ func NewYQLAgent(cfgen *ytconfig.Generator, ytsaurus *apiproxy.Ytsaurus, master
 		func() ([]byte, error) {
 			return cfgen.GetYQLAgentConfig(resource.Spec.YQLAgents)
 		},
+		WithContainerPorts(corev1.ContainerPort{
+			Name:          consts.YTRPCPortName,
+			ContainerPort: consts.YQLAgentRPCPort,
+			Protocol:      corev1.ProtocolTCP,
+		}),
 	)
 
 	return &YqlAgent{

diff --git a/pkg/consts/address.go b/pkg/consts/address.go
@@ -1,7 +1,12 @@
 package consts
 
-const YTMonitoringPortName = "ytsaurus-metrics"
-const YTMonitoringPort = 10000
+const (
+	YTRPCPortName = "rpc"
+
+	YTMonitoringContainerPortName = "metrics"
+	YTMonitoringServicePortName   = "ytsaurus-metrics"
+	YTMonitoringPort              = 10000
+)
 
 const (
 	DiscoveryRPCPort        = 9020