From 7de579a236d287175746dcd70c561adcc1a1d305 Mon Sep 17 00:00:00 2001
From: Matthieu Vachon <matt@streamingfast.io>
Date: Thu, 18 Jan 2024 23:58:51 -0500
Subject: [PATCH] Substreams app(s) gained ability to check if shutdown is
 pending

---
 app/tier1.go |  7 ++++++-
 app/tier2.go | 14 ++++++++++++--
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/app/tier1.go b/app/tier1.go
index 9b253264e..be908bf3b 100644
--- a/app/tier1.go
+++ b/app/tier1.go
@@ -28,6 +28,7 @@ type Tier1Modules struct {
 	Authenticator         dauth.Authenticator
 	HeadTimeDriftMetric   *dmetrics.HeadTimeDrift
 	HeadBlockNumberMetric *dmetrics.HeadBlockNum
+	CheckPendingShutDown  func() bool
 }
 
 type Tier1Config struct {
@@ -191,7 +192,7 @@ func (a *Tier1App) Run() error {
 		}
 
 		a.logger.Info("launching gRPC server", zap.Bool("live_support", withLive))
-		a.isReady.CAS(false, true)
+		a.isReady.CompareAndSwap(false, true)
 
 		err := service.ListenTier1(a.config.GRPCListenAddr, svc, a.modules.Authenticator, a.logger, a.HealthCheck)
 		a.Shutdown(err)
@@ -214,6 +215,10 @@ func (a *Tier1App) IsReady(ctx context.Context) bool {
 		return false
 	}
 
+	if a.modules.CheckPendingShutDown != nil && a.modules.CheckPendingShutDown() {
+		return false
+	}
+
 	return a.isReady.Load()
 }
 
diff --git a/app/tier2.go b/app/tier2.go
index dddde432a..e8ab3ef46 100644
--- a/app/tier2.go
+++ b/app/tier2.go
@@ -35,14 +35,20 @@ type Tier2Config struct {
 type Tier2App struct {
 	*shutter.Shutter
 	config  *Tier2Config
+	modules *Tier2Modules
 	logger  *zap.Logger
 	isReady *atomic.Bool
 }
 
-func NewTier2(logger *zap.Logger, config *Tier2Config) *Tier2App {
+type Tier2Modules struct {
+	CheckPendingShutDown func() bool
+}
+
+func NewTier2(logger *zap.Logger, config *Tier2Config, modules *Tier2Modules) *Tier2App {
 	return &Tier2App{
 		Shutter: shutter.New(),
 		config:  config,
+		modules: modules,
 		logger:  logger,
 
 		isReady: atomic.NewBool(false),
@@ -100,7 +106,7 @@ func (a *Tier2App) Run() error {
 
 	go func() {
 		a.logger.Info("launching gRPC server")
-		a.isReady.CAS(false, true)
+		a.isReady.CompareAndSwap(false, true)
 
 		err := service.ListenTier2(a.config.GRPCListenAddr, a.config.ServiceDiscoveryURL, svc, trustAuth, a.logger, a.HealthCheck)
 		a.Shutdown(err)
@@ -120,6 +126,10 @@ func (a *Tier2App) IsReady(ctx context.Context) bool {
 		return false
 	}
 
+	if a.modules.CheckPendingShutDown != nil && a.modules.CheckPendingShutDown() {
+		return false
+	}
+
 	return a.isReady.Load()
 }