From 9e928ad4644c68cdc4ac5426bc10895100bf30a7 Mon Sep 17 00:00:00 2001
From: maskpp <maskpp266@gmail.com>
Date: Tue, 2 Jul 2024 22:56:10 +0800
Subject: [PATCH] chore(relayer): close db connection before stop service.
 (#17709)

---
 packages/relayer/event.go                     |  1 +
 packages/relayer/indexer/indexer.go           |  5 +++++
 packages/relayer/pkg/http/server.go           |  6 ++++++
 packages/relayer/pkg/mock/event_repository.go |  5 +++++
 packages/relayer/pkg/repo/event.go            | 10 ++++++++++
 packages/relayer/processor/processor.go       |  5 +++++
 packages/relayer/watchdog/watchdog.go         |  5 +++++
 7 files changed, 37 insertions(+)

diff --git a/packages/relayer/event.go b/packages/relayer/event.go
index f2e8c9badcd..cc2afc3558f 100644
--- a/packages/relayer/event.go
+++ b/packages/relayer/event.go
@@ -128,6 +128,7 @@ type FindAllByAddressOpts struct {
 
 // EventRepository is used to interact with events in the store
 type EventRepository interface {
+	Close() error
 	Save(ctx context.Context, opts *SaveEventOpts) (*Event, error)
 	UpdateStatus(ctx context.Context, id int, status EventStatus) error
 	UpdateFeesAndProfitability(ctx context.Context, id int, opts *UpdateFeesAndProfitabilityOpts) error
diff --git a/packages/relayer/indexer/indexer.go b/packages/relayer/indexer/indexer.go
index 80d0b330a90..55a05dd5d23 100644
--- a/packages/relayer/indexer/indexer.go
+++ b/packages/relayer/indexer/indexer.go
@@ -252,6 +252,11 @@ func (i *Indexer) Name() string {
 // context is stopped externally by cmd/main.go shutdown.
 func (i *Indexer) Close(ctx context.Context) {
 	i.wg.Wait()
+
+	// Close db connection.
+	if err := i.eventRepo.Close(); err != nil {
+		slog.Error("Failed to close db connection", "err", err)
+	}
 }
 
 // Start starts the indexer, which should initialize the queue, add to wait groups,
diff --git a/packages/relayer/pkg/http/server.go b/packages/relayer/pkg/http/server.go
index 9439e711f92..c3739197f9f 100644
--- a/packages/relayer/pkg/http/server.go
+++ b/packages/relayer/pkg/http/server.go
@@ -2,6 +2,7 @@ package http
 
 import (
 	"context"
+	"log/slog"
 	"math/big"
 	"net/http"
 	"os"
@@ -130,6 +131,11 @@ func (srv *Server) Start(address string) error {
 
 // Shutdown shuts down the HTTP server
 func (srv *Server) Shutdown(ctx context.Context) error {
+	// Close db connection.
+	if err := srv.eventRepo.Close(); err != nil {
+		slog.Error("Failed to close db connection", "err", err)
+	}
+
 	return srv.echo.Shutdown(ctx)
 }
 
diff --git a/packages/relayer/pkg/mock/event_repository.go b/packages/relayer/pkg/mock/event_repository.go
index 29665c393b5..0456ba7c2c6 100644
--- a/packages/relayer/pkg/mock/event_repository.go
+++ b/packages/relayer/pkg/mock/event_repository.go
@@ -22,6 +22,11 @@ func NewEventRepository() *EventRepository {
 		events: make([]*relayer.Event, 0),
 	}
 }
+
+func (r *EventRepository) Close() error {
+	return nil
+}
+
 func (r *EventRepository) Save(ctx context.Context, opts *relayer.SaveEventOpts) (*relayer.Event, error) {
 	r.events = append(r.events, &relayer.Event{
 		ID:           rand.Int(), // nolint: gosec
diff --git a/packages/relayer/pkg/repo/event.go b/packages/relayer/pkg/repo/event.go
index 80d9e149d09..3309c9982f8 100644
--- a/packages/relayer/pkg/repo/event.go
+++ b/packages/relayer/pkg/repo/event.go
@@ -30,6 +30,16 @@ func NewEventRepository(dbHandler db.DB) (*EventRepository, error) {
 	}, nil
 }
 
+// Close closes the database connection.
+func (r *EventRepository) Close() error {
+	sqlDB, err := r.db.DB()
+	if err != nil {
+		return err
+	}
+
+	return sqlDB.Close()
+}
+
 func (r *EventRepository) Save(ctx context.Context, opts *relayer.SaveEventOpts) (*relayer.Event, error) {
 	e := &relayer.Event{
 		Data:                   datatypes.JSON(opts.Data),
diff --git a/packages/relayer/processor/processor.go b/packages/relayer/processor/processor.go
index a618d8060b8..5e537a069a2 100644
--- a/packages/relayer/processor/processor.go
+++ b/packages/relayer/processor/processor.go
@@ -383,6 +383,11 @@ func (p *Processor) Close(ctx context.Context) {
 	p.cancel()
 
 	p.wg.Wait()
+
+	// Close db connection.
+	if err := p.eventRepo.Close(); err != nil {
+		slog.Error("Failed to close db connection", "err", err)
+	}
 }
 
 func (p *Processor) Start() error {
diff --git a/packages/relayer/watchdog/watchdog.go b/packages/relayer/watchdog/watchdog.go
index 8c4f0eb329c..36ed241793f 100644
--- a/packages/relayer/watchdog/watchdog.go
+++ b/packages/relayer/watchdog/watchdog.go
@@ -202,6 +202,11 @@ func (w *Watchdog) Close(ctx context.Context) {
 	w.cancel()
 
 	w.wg.Wait()
+
+	// Close db connection.
+	if err := w.eventRepo.Close(); err != nil {
+		slog.Error("Failed to close db connection", "err", err)
+	}
 }
 
 func (w *Watchdog) Start() error {