Skip to content

Commit a453319

Browse files
cre8ivejpcre8ivejp
committed
chore: change api agteway server context
Signed-off-by: Alessandro Yuichi Okimoto <yuichijpn@gmail.com>
1 parent 1d75ac9 commit a453319

10 files changed

Lines changed: 154 additions & 87 deletions

File tree

‎manifests/bucketeer/charts/api/templates/deployment.yaml‎

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -189,11 +189,32 @@ spec:
189189
190190
# Wait for active requests to drain
191191
for i in $(seq 1 "$max_wait"); do
192-
active=$(wget -q -O- "http://localhost:${admin_port}/stats" 2>/dev/null | grep "http.ingress_http.downstream_rq_active" | awk '{print $2}' || echo "0")
193-
[ -z "$active" ] && active=0
194-
[ "$active" -eq 0 ] && break
192+
# Fetch stats and check if request succeeded
193+
stats=$(wget -q -T 1 -O- "http://127.0.0.1:${admin_port}/stats" 2>/dev/null)
194+
if [ $? -ne 0 ] || [ -z "$stats" ]; then
195+
echo "Check $i/$max_wait: Failed to fetch stats, retrying..."
196+
sleep 1
197+
continue
198+
fi
199+
200+
# Extract active requests metric
201+
active=$(echo "$stats" | grep -E '^http\.ingress_http\.downstream_rq_active:' | awk '{print $2}')
202+
if [ -z "$active" ]; then
203+
echo "Check $i/$max_wait: Metric not found, retrying..."
204+
sleep 1
205+
continue
206+
fi
207+
208+
echo "Check $i/$max_wait: Active requests: $active"
209+
[ "$active" -eq 0 ] && echo "No active requests, exiting gracefully" && break
195210
sleep 1
196211
done
212+
213+
if [ -n "$active" ] && [ "$active" -eq 0 ]; then
214+
echo "Graceful shutdown completed successfully"
215+
else
216+
echo "Warning: Timed out after ${max_wait}s - active=${active:-unknown}"
217+
fi
197218
exit 0
198219
command: ["envoy"]
199220
args:

‎manifests/bucketeer/charts/batch/templates/deployment.yaml‎

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -223,19 +223,36 @@ spec:
223223
- -c
224224
- |
225225
admin_port={{ .Values.envoy.adminPort }}
226-
max_wait=35
227-
propagation_delay=15
228-
229-
# Wait for GCLB to detect unhealthy status and stop routing
230-
sleep "$propagation_delay"
226+
max_wait=60
231227
232228
# Wait for active requests to drain
233229
for i in $(seq 1 "$max_wait"); do
234-
active=$(wget -q -O- "http://localhost:${admin_port}/stats" 2>/dev/null | grep "http.ingress_http.downstream_rq_active" | awk '{print $2}' || echo "0")
235-
[ -z "$active" ] && active=0
236-
[ "$active" -eq 0 ] && break
230+
# Fetch stats and check if request succeeded
231+
stats=$(wget -q -T 1 -O- "http://127.0.0.1:${admin_port}/stats" 2>/dev/null)
232+
if [ $? -ne 0 ] || [ -z "$stats" ]; then
233+
echo "Check $i/$max_wait: Failed to fetch stats, retrying..."
234+
sleep 1
235+
continue
236+
fi
237+
238+
# Extract active requests metric
239+
active=$(echo "$stats" | grep -E '^http\.ingress_http\.downstream_rq_active:' | awk '{print $2}')
240+
if [ -z "$active" ]; then
241+
echo "Check $i/$max_wait: Metric not found, retrying..."
242+
sleep 1
243+
continue
244+
fi
245+
246+
echo "Check $i/$max_wait: Active requests: $active"
247+
[ "$active" -eq 0 ] && echo "No active requests, exiting gracefully" && break
237248
sleep 1
238249
done
250+
251+
if [ -n "$active" ] && [ "$active" -eq 0 ]; then
252+
echo "Graceful shutdown completed successfully"
253+
else
254+
echo "Warning: Timed out after ${max_wait}s - active=${active:-unknown}"
255+
fi
239256
exit 0
240257
command: ["envoy"]
241258
args:

‎manifests/bucketeer/charts/subscriber/templates/deployment.yaml‎

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -203,19 +203,36 @@ spec:
203203
- -c
204204
- |
205205
admin_port={{ .Values.envoy.adminPort }}
206-
max_wait=35
207-
propagation_delay=15
208-
209-
# Wait for GCLB to detect unhealthy status and stop routing
210-
sleep "$propagation_delay"
206+
max_wait=60
211207
212208
# Wait for active requests to drain
213209
for i in $(seq 1 "$max_wait"); do
214-
active=$(wget -q -O- "http://localhost:${admin_port}/stats" 2>/dev/null | grep "http.ingress_http.downstream_rq_active" | awk '{print $2}' || echo "0")
215-
[ -z "$active" ] && active=0
216-
[ "$active" -eq 0 ] && break
210+
# Fetch stats and check if request succeeded
211+
stats=$(wget -q -T 1 -O- "http://127.0.0.1:${admin_port}/stats" 2>/dev/null)
212+
if [ $? -ne 0 ] || [ -z "$stats" ]; then
213+
echo "Check $i/$max_wait: Failed to fetch stats, retrying..."
214+
sleep 1
215+
continue
216+
fi
217+
218+
# Extract active requests metric
219+
active=$(echo "$stats" | grep -E '^http\.ingress_http\.downstream_rq_active:' | awk '{print $2}')
220+
if [ -z "$active" ]; then
221+
echo "Check $i/$max_wait: Metric not found, retrying..."
222+
sleep 1
223+
continue
224+
fi
225+
226+
echo "Check $i/$max_wait: Active requests: $active"
227+
[ "$active" -eq 0 ] && echo "No active requests, exiting gracefully" && break
217228
sleep 1
218229
done
230+
231+
if [ -n "$active" ] && [ "$active" -eq 0 ]; then
232+
echo "Graceful shutdown completed successfully"
233+
else
234+
echo "Warning: Timed out after ${max_wait}s - active=${active:-unknown}"
235+
fi
219236
exit 0
220237
command: ["envoy"]
221238
args:

‎manifests/bucketeer/charts/web/templates/deployment.yaml‎

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -273,19 +273,36 @@ spec:
273273
- -c
274274
- |
275275
admin_port={{ .Values.envoy.adminPort }}
276-
max_wait=35
277-
propagation_delay=15
278-
279-
# Wait for GCLB to detect unhealthy status and stop routing
280-
sleep "$propagation_delay"
276+
max_wait=60
281277
282278
# Wait for active requests to drain
283279
for i in $(seq 1 "$max_wait"); do
284-
active=$(wget -q -O- "http://localhost:${admin_port}/stats" 2>/dev/null | grep "http.ingress_http.downstream_rq_active" | awk '{print $2}' || echo "0")
285-
[ -z "$active" ] && active=0
286-
[ "$active" -eq 0 ] && break
280+
# Fetch stats and check if request succeeded
281+
stats=$(wget -q -T 1 -O- "http://127.0.0.1:${admin_port}/stats" 2>/dev/null)
282+
if [ $? -ne 0 ] || [ -z "$stats" ]; then
283+
echo "Check $i/$max_wait: Failed to fetch stats, retrying..."
284+
sleep 1
285+
continue
286+
fi
287+
288+
# Extract active requests metric
289+
active=$(echo "$stats" | grep -E '^http\.ingress_http\.downstream_rq_active:' | awk '{print $2}')
290+
if [ -z "$active" ]; then
291+
echo "Check $i/$max_wait: Metric not found, retrying..."
292+
sleep 1
293+
continue
294+
fi
295+
296+
echo "Check $i/$max_wait: Active requests: $active"
297+
[ "$active" -eq 0 ] && echo "No active requests, exiting gracefully" && break
287298
sleep 1
288299
done
300+
301+
if [ -n "$active" ] && [ "$active" -eq 0 ]; then
302+
echo "Graceful shutdown completed successfully"
303+
else
304+
echo "Warning: Timed out after ${max_wait}s - active=${active:-unknown}"
305+
fi
289306
exit 0
290307
command: ["envoy"]
291308
args:

‎pkg/api/cmd/server.go‎

Lines changed: 16 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -230,9 +230,6 @@ func RegisterCommand(r cli.CommandRegistry, p cli.ParentCommand) cli.Command {
230230
func (s *server) Run(ctx context.Context, metrics metrics.Metrics, logger *zap.Logger) error {
231231
registerer := metrics.DefaultRegisterer()
232232

233-
pubsubCtx, pubsubCancel := context.WithTimeout(ctx, 5*time.Second)
234-
defer pubsubCancel()
235-
236233
// Create PubSub client using the factory
237234
pubSubType := factory.PubSubType(*s.pubSubType)
238235
factoryOpts := []factory.Option{
@@ -261,6 +258,8 @@ func (s *server) Run(ctx context.Context, metrics metrics.Metrics, logger *zap.L
261258
factoryOpts = append(factoryOpts, factory.WithPartitionCount(*s.pubSubRedisPartitionCount))
262259
}
263260

261+
pubsubCtx, pubsubCancel := context.WithCancel(context.Background())
262+
defer pubsubCancel()
264263
pubsubClient, err := factory.NewClient(pubsubCtx, factoryOpts...)
265264
if err != nil {
266265
return err
@@ -388,7 +387,7 @@ func (s *server) Run(ctx context.Context, metrics metrics.Metrics, logger *zap.L
388387
if err != nil {
389388
return err
390389
}
391-
defer auditLogClient.Close()
390+
defer autoOpsClient.Close()
392391

393392
tagClient, err := tagclient.NewClient(*s.tagService, *s.certPath,
394393
client.WithPerRPCCredentials(creds),
@@ -502,8 +501,8 @@ func (s *server) Run(ctx context.Context, metrics metrics.Metrics, logger *zap.L
502501
// We don't check the Redis health status because if the check fails,
503502
// the Kubernetes will restart the container and it might cause internal errors.
504503
// Use a dedicated context so we can stop the health checker goroutine cleanly during shutdown
505-
healthCheckCtx, healthCheckCancel := context.WithCancel(ctx)
506-
defer healthCheckCancel() // Ensure cleanup on all paths (including early returns)
504+
healthCheckCtx, healthCheckCancel := context.WithCancel(context.Background())
505+
defer healthCheckCancel()
507506

508507
healthChecker := health.NewGrpcChecker(
509508
health.WithTimeout(5*time.Second),
@@ -545,7 +544,9 @@ func (s *server) Run(ctx context.Context, metrics metrics.Metrics, logger *zap.L
545544
return fmt.Errorf("failed to create API gateway: %v", err)
546545
}
547546

548-
if err := apiGateway.Start(ctx, gatewayHandler); err != nil {
547+
serverCtx, serverCtxCancel := context.WithCancel(context.Background())
548+
defer serverCtxCancel()
549+
if err := apiGateway.Start(serverCtx, gatewayHandler); err != nil {
549550
return fmt.Errorf("failed to start API gateway: %v", err)
550551
}
551552

@@ -584,27 +585,20 @@ func (s *server) Run(ctx context.Context, metrics metrics.Metrics, logger *zap.L
584585
shutdownStartTime := time.Now()
585586
logger.Info("Starting graceful shutdown sequence")
586587

587-
waitBeforeUnready := 10 * time.Second
588-
logger.Info("Waiting before marking unready",
589-
zap.Duration("wait_before_unready", waitBeforeUnready))
590-
time.Sleep(waitBeforeUnready)
588+
// Wait for K8s endpoint propagation
589+
// This prevents "context deadline exceeded" errors during high traffic.
590+
time.Sleep(propagationDelay)
591+
logger.Info("Starting HTTP/gRPC server shutdown")
591592

592-
// Cancel the health checker goroutines to prevent connection errors during shutdown
593-
healthCheckCancel()
594593
// Mark as unhealthy so readiness probes fail
595594
// This ensures Kubernetes readiness probe fails on next check,
596595
// preventing new traffic from being routed to this pod.
597596
healthChecker.Stop()
598597
restHealthChecker.Stop()
599598

600-
// Wait for K8s endpoint propagation
601-
// This prevents "context deadline exceeded" errors during high traffic.
602-
time.Sleep(propagationDelay)
603-
logger.Info("Starting HTTP/gRPC server shutdown")
604-
605-
// CRITICAL: Shutdown order matters due to dependencies:
599+
// Shutdown order matters due to dependencies:
606600
// 1. apiGateway/httpServer make gRPC calls to the backend server
607-
// 2. We MUST drain them BEFORE stopping the backend
601+
// 2. We MUST drain them BEFORE stopping the backend sever
608602
// 3. Otherwise their handlers hang waiting for a dead backend
609603
// We run apiGateway and httpServer in parallel since they don't depend on each other
610604
var wg sync.WaitGroup
@@ -623,9 +617,11 @@ func (s *server) Run(ctx context.Context, metrics metrics.Metrics, logger *zap.L
623617

624618
// Wait for HTTP/REST traffic to fully drain
625619
wg.Wait()
620+
logger.Info("gRPC-gateway and HTTP server shutdown completed")
626621

627622
// Now it's safe to stop the gRPC server (no more HTTP→gRPC calls)
628623
server.Stop(grpcStopTimeout)
624+
logger.Info("gRPC server shutdown completed")
629625

630626
// Close clients
631627
// These are fast cleanup operations that can run asynchronously.

‎pkg/batch/cmd/server/server.go‎

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -584,8 +584,8 @@ func (s *server) Run(ctx context.Context, metrics metrics.Metrics, logger *zap.L
584584
)
585585

586586
// Use a dedicated context so we can stop the health checker goroutine cleanly during shutdown
587-
healthCheckCtx, healthCheckCancel := context.WithCancel(ctx)
588-
defer healthCheckCancel() // Ensure cleanup on all paths (including early returns)
587+
healthCheckCtx, healthCheckCancel := context.WithCancel(context.Background())
588+
defer healthCheckCancel()
589589

590590
healthChecker := health.NewGrpcChecker(
591591
health.WithTimeout(time.Second),
@@ -627,31 +627,33 @@ func (s *server) Run(ctx context.Context, metrics metrics.Metrics, logger *zap.L
627627
return fmt.Errorf("failed to create batch gateway: %v", err)
628628
}
629629

630-
if err := batchGateway.Start(ctx, batchHandler); err != nil {
630+
batchCtx, batchCancel := context.WithCancel(context.Background())
631+
defer batchCancel()
632+
if err := batchGateway.Start(batchCtx, batchHandler); err != nil {
631633
return fmt.Errorf("failed to start batch gateway: %v", err)
632634
}
633635

634636
defer func() {
635637
shutdownStartTime := time.Now()
636638
logger.Info("Starting graceful shutdown sequence")
637639

638-
// Cancel the health checker goroutines to prevent connection errors during shutdown
639-
healthCheckCancel()
640-
// Mark as unhealthy so readiness probes fail
641-
// This ensures Kubernetes readiness probe fails on next check,
642-
// preventing new traffic from being routed to this pod.
643-
healthChecker.Stop()
644-
645640
// Wait for K8s endpoint propagation
646641
// This prevents "context deadline exceeded" errors during high traffic.
647642
time.Sleep(propagationDelay)
648643
logger.Info("Starting HTTP/gRPC server shutdown")
649644

650-
// Gracefully stop REST gateway (calls the gRPC server internally)
645+
// Mark as unhealthy so readiness probes fail
646+
// This ensures Kubernetes readiness probe fails on next check,
647+
// preventing new traffic from being routed to this pod.
648+
healthChecker.Stop()
649+
650+
// Gracefully stop gRPC Gateway (calls the gRPC server internally)
651651
batchGateway.Stop(serverShutDownTimeout)
652+
logger.Info("gRPC-gateway server shutdown completed")
652653

653654
// Stop gRPC server (only pure gRPC connections remain)
654655
server.Stop(grpcStopTimeout)
656+
logger.Info("gRPC server shutdown completed")
655657

656658
// Close clients
657659
// These are fast cleanup operations that can run asynchronously.

‎pkg/subscriber/cmd/server/server.go‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -343,7 +343,7 @@ func (s *server) Run(ctx context.Context, metrics metrics.Metrics, logger *zap.L
343343
// healthCheckService
344344
// Use a dedicated context so we can stop the health checker goroutine cleanly during shutdown
345345
healthCheckCtx, healthCheckCancel := context.WithCancel(ctx)
346-
defer healthCheckCancel() // Ensure cleanup on all paths (including early returns)
346+
defer healthCheckCancel()
347347

348348
restHealthChecker := health.NewRestChecker(
349349
"", "",
@@ -365,8 +365,6 @@ func (s *server) Run(ctx context.Context, metrics metrics.Metrics, logger *zap.L
365365
shutdownStartTime := time.Now()
366366
logger.Info("Starting graceful shutdown sequence")
367367

368-
// Cancel the health checker goroutines to prevent connection errors during shutdown
369-
healthCheckCancel()
370368
// Mark as unhealthy so readiness probes fail
371369
// This ensures Kubernetes readiness probe fails on next check,
372370
// preventing new traffic from being routed to this pod.
@@ -376,6 +374,7 @@ func (s *server) Run(ctx context.Context, metrics metrics.Metrics, logger *zap.L
376374
// Stop PubSub subscription
377375
// This stops receiving new messages and allows in-flight messages to be processed.
378376
multiPubSub.Stop()
377+
logger.Info("PubSub subscription stopped, all messages processed")
379378

380379
// Close clients
381380
// These are fast cleanup operations that can run asynchronously.

‎pkg/subscriber/on_demand_subscriber.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -233,8 +233,9 @@ func (s *onDemandSubscriber) createPubSubClient(ctx context.Context) error {
233233
}
234234
}
235235

236-
// Create the PubSub client using the factory
237-
pubsubClient, err := factory.NewClient(ctx, factoryOpts...)
236+
// Create the PubSub client using the factory with context.Background()
237+
// to ensure connections remain healthy until explicitly stopped during graceful shutdown
238+
pubsubClient, err := factory.NewClient(context.Background(), factoryOpts...)
238239
if err != nil {
239240
s.logger.Error("Failed to create pubsub client",
240241
zap.Error(err),

‎pkg/subscriber/subscriber.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -201,8 +201,9 @@ func (s pubSubSubscriber) createPuller(
201201
}
202202
}
203203

204-
// Create the PubSub client using the factory
205-
pubsubClient, err = factory.NewClient(ctx, factoryOpts...)
204+
// Create the PubSub client using the factory with context.Background()
205+
// to ensure connections remain healthy until explicitly stopped during graceful shutdown
206+
pubsubClient, err = factory.NewClient(context.Background(), factoryOpts...)
206207
if err != nil {
207208
s.logger.Error("Failed to create pubsub client",
208209
zap.Error(err),

0 commit comments

Comments
 (0)