@@ -37,11 +37,19 @@ type K8sCircuitBreakHandler struct {
3737 client kubernetes.Interface
3838 defaultNamespace string
3939 now func () time.Time
40+ afterFunc func (time.Duration , func ()) * time.Timer
4041
4142 mu sync.Mutex
4243 timers map [string ]* time.Timer
4344}
4445
46+ const (
47+ circuitBreakPolicyPrefix = "circuit-break-"
48+ circuitBreakExpiresAtAnnotation = "underpass.ai/expires_at"
49+ circuitBreakTargetServiceAnnotation = "underpass.ai/target_service"
50+ circuitBreakDownstreamAnnotation = "underpass.ai/downstream"
51+ )
52+
4553func NewK8sScaleDeploymentHandler (client kubernetes.Interface , defaultNamespace string ) * K8sScaleDeploymentHandler {
4654 return & K8sScaleDeploymentHandler {client : client , defaultNamespace : strings .TrimSpace (defaultNamespace )}
4755}
@@ -51,12 +59,66 @@ func NewK8sRestartPodsHandler(client kubernetes.Interface, defaultNamespace stri
5159}
5260
5361func NewK8sCircuitBreakHandler (client kubernetes.Interface , defaultNamespace string ) * K8sCircuitBreakHandler {
54- return & K8sCircuitBreakHandler {
62+ return newK8sCircuitBreakHandler (client , defaultNamespace , nil , nil )
63+ }
64+
65+ func newK8sCircuitBreakHandler (
66+ client kubernetes.Interface ,
67+ defaultNamespace string ,
68+ now func () time.Time ,
69+ afterFunc func (time.Duration , func ()) * time.Timer ,
70+ ) * K8sCircuitBreakHandler {
71+ if now == nil {
72+ now = func () time.Time { return time .Now ().UTC () }
73+ }
74+ if afterFunc == nil {
75+ afterFunc = time .AfterFunc
76+ }
77+
78+ handler := & K8sCircuitBreakHandler {
5579 client : client ,
5680 defaultNamespace : strings .TrimSpace (defaultNamespace ),
57- now : func () time.Time { return time .Now ().UTC () },
81+ now : now ,
82+ afterFunc : afterFunc ,
5883 timers : map [string ]* time.Timer {},
5984 }
85+ handler .reconcileExistingPolicies (context .Background ())
86+ return handler
87+ }
88+
89+ func (h * K8sCircuitBreakHandler ) reconcileExistingPolicies (ctx context.Context ) {
90+ if h == nil || h .client == nil {
91+ return
92+ }
93+
94+ namespace := strings .TrimSpace (h .defaultNamespace )
95+ policies , err := h .client .NetworkingV1 ().NetworkPolicies (namespace ).List (ctx , metav1.ListOptions {})
96+ if err != nil {
97+ return
98+ }
99+
100+ now := h .now ().UTC ()
101+ for i := range policies .Items {
102+ policy := policies .Items [i ]
103+ if ! strings .HasPrefix (policy .Name , circuitBreakPolicyPrefix ) {
104+ continue
105+ }
106+ expiresAtRaw := strings .TrimSpace (policy .Annotations [circuitBreakExpiresAtAnnotation ])
107+ if expiresAtRaw == "" {
108+ continue
109+ }
110+ expiresAt , parseErr := time .Parse (time .RFC3339 , expiresAtRaw )
111+ if parseErr != nil {
112+ continue
113+ }
114+
115+ ttl := expiresAt .Sub (now )
116+ if ttl <= 0 {
117+ _ = h .client .NetworkingV1 ().NetworkPolicies (policy .Namespace ).Delete (ctx , policy .Name , metav1.DeleteOptions {})
118+ continue
119+ }
120+ h .schedulePolicyCleanup (policy .Namespace , policy .Name , ttl )
121+ }
60122}
61123
62124func (h * K8sScaleDeploymentHandler ) Name () string {
@@ -166,7 +228,7 @@ func (h *K8sRestartPodsHandler) Invoke(ctx context.Context, session domain.Sessi
166228
167229 mode := strings .TrimSpace (request .Mode )
168230 if mode == "" {
169- mode = "rollout_restart"
231+ return app. ToolRunResult {}, k8sInvalidArgument ( " mode is required" )
170232 }
171233 output := map [string ]any {
172234 k8sDelivKeyNamespace : namespace ,
@@ -324,7 +386,7 @@ func buildRestartPodsSelector(deployment *appsv1.Deployment, labelSelector strin
324386
325387func circuitBreakPolicyID (namespace , targetService , downstream string ) string {
326388 hash := sha256 .Sum256 ([]byte (namespace + ":" + targetService + ":" + downstream ))
327- return fmt .Sprintf ("circuit-break-%x" , hash [:6 ])
389+ return fmt .Sprintf ("%s%x" , circuitBreakPolicyPrefix , hash [:6 ])
328390}
329391
330392func buildCircuitBreakNetworkPolicy (
@@ -338,9 +400,9 @@ func buildCircuitBreakNetworkPolicy(
338400 Name : policyID ,
339401 Namespace : namespace ,
340402 Annotations : map [string ]string {
341- "underpass.ai/target_service" : targetService .Name ,
342- "underpass.ai/downstream" : downstreamService .Name ,
343- "underpass.ai/expires_at" : expiresAt .UTC ().Format (time .RFC3339 ),
403+ circuitBreakTargetServiceAnnotation : targetService .Name ,
404+ circuitBreakDownstreamAnnotation : downstreamService .Name ,
405+ circuitBreakExpiresAtAnnotation : expiresAt .UTC ().Format (time .RFC3339 ),
344406 },
345407 },
346408 Spec : networkingv1.NetworkPolicySpec {
@@ -380,7 +442,7 @@ func (h *K8sCircuitBreakHandler) schedulePolicyCleanup(namespace, policyID strin
380442 if existing := h .timers [policyID ]; existing != nil {
381443 existing .Stop ()
382444 }
383- timer := time . AfterFunc (ttl , func () {
445+ timer := h . afterFunc (ttl , func () {
384446 ctx , cancel := context .WithTimeout (context .Background (), 30 * time .Second )
385447 defer cancel ()
386448 _ = h .client .NetworkingV1 ().NetworkPolicies (namespace ).Delete (ctx , policyID , metav1.DeleteOptions {})
0 commit comments