Skip to content

Commit 033732a

Browse files
committed
Add timeouts to yugabytedb operations
1 parent 37cae32 commit 033732a

3 files changed

Lines changed: 36 additions & 8 deletions

File tree

api/pkg/di/container.go

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -239,14 +239,14 @@ func (container *Container) YugaByteDB() (db *gorm.DB) {
239239
sql.SetMaxIdleConns(3)
240240
sql.SetConnMaxLifetime(time.Minute * 10)
241241

242-
container.logger.Debug(fmt.Sprintf("Running migrations for yugabyte [%T]", db))
243-
if err = db.AutoMigrate(&entities.Heartbeat{}); err != nil {
244-
container.logger.Fatal(stacktrace.Propagate(err, fmt.Sprintf("cannot migrate %T", &entities.Heartbeat{})))
245-
}
246-
247-
if err = db.AutoMigrate(&entities.HeartbeatMonitor{}); err != nil {
248-
container.logger.Fatal(stacktrace.Propagate(err, fmt.Sprintf("cannot migrate %T", &entities.HeartbeatMonitor{})))
249-
}
242+
//container.logger.Debug(fmt.Sprintf("Running migrations for yugabyte [%T]", db))
243+
//if err = db.AutoMigrate(&entities.Heartbeat{}); err != nil {
244+
// container.logger.Fatal(stacktrace.Propagate(err, fmt.Sprintf("cannot migrate %T", &entities.Heartbeat{})))
245+
//}
246+
//
247+
//if err = db.AutoMigrate(&entities.HeartbeatMonitor{}); err != nil {
248+
// container.logger.Fatal(stacktrace.Propagate(err, fmt.Sprintf("cannot migrate %T", &entities.HeartbeatMonitor{})))
249+
//}
250250

251251
container.yugaByteDB = db
252252
return container.yugaByteDB

api/pkg/repositories/gorm_heartbeat_monitor_repository.go

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,9 @@ func (repository *gormHeartbeatMonitorRepository) UpdateQueueID(ctx context.Cont
2727
ctx, span := repository.tracer.Start(ctx)
2828
defer span.End()
2929

30+
ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
31+
defer cancel()
32+
3033
err := repository.db.
3134
Model(&entities.HeartbeatMonitor{}).
3235
Where("id = ?", monitorID).
@@ -45,6 +48,9 @@ func (repository *gormHeartbeatMonitorRepository) Delete(ctx context.Context, us
4548
ctx, span := repository.tracer.Start(ctx)
4649
defer span.End()
4750

51+
ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
52+
defer cancel()
53+
4854
err := repository.db.WithContext(ctx).
4955
Where("user_id = ?", userID).
5056
Where("owner = ?", owner).
@@ -75,6 +81,9 @@ func (repository *gormHeartbeatMonitorRepository) Index(ctx context.Context, use
7581
ctx, span := repository.tracer.Start(ctx)
7682
defer span.End()
7783

84+
ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
85+
defer cancel()
86+
7887
query := repository.db.WithContext(ctx).Where("user_id = ?", userID).Where("owner = ?", owner)
7988
heartbeats := new([]entities.Heartbeat)
8089
if err := query.Order("timestamp DESC").Limit(params.Limit).Offset(params.Skip).Find(&heartbeats).Error; err != nil {
@@ -90,6 +99,9 @@ func (repository *gormHeartbeatMonitorRepository) Store(ctx context.Context, hea
9099
ctx, span := repository.tracer.Start(ctx)
91100
defer span.End()
92101

102+
ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
103+
defer cancel()
104+
93105
if err := repository.db.WithContext(ctx).Create(heartbeatMonitor).Error; err != nil {
94106
msg := fmt.Sprintf("cannot save heartbeatMonitor monitor with ID [%s]", heartbeatMonitor.ID)
95107
return repository.tracer.WrapErrorSpan(span, stacktrace.Propagate(err, msg))
@@ -103,6 +115,9 @@ func (repository *gormHeartbeatMonitorRepository) Load(ctx context.Context, user
103115
ctx, span := repository.tracer.Start(ctx)
104116
defer span.End()
105117

118+
ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
119+
defer cancel()
120+
106121
phone := new(entities.HeartbeatMonitor)
107122
err := repository.db.WithContext(ctx).
108123
Where("user_id = ?", userID).
@@ -127,6 +142,9 @@ func (repository *gormHeartbeatMonitorRepository) Exists(ctx context.Context, us
127142
ctx, span := repository.tracer.Start(ctx)
128143
defer span.End()
129144

145+
ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
146+
defer cancel()
147+
130148
var exists bool
131149
err := repository.db.WithContext(ctx).
132150
Model(&entities.HeartbeatMonitor{}).

api/pkg/repositories/gorm_heartbeat_repository.go

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package repositories
33
import (
44
"context"
55
"fmt"
6+
"time"
67

78
"github.com/pkg/errors"
89

@@ -36,6 +37,9 @@ func (repository *gormHeartbeatRepository) Last(ctx context.Context, userID enti
3637
ctx, span := repository.tracer.Start(ctx)
3738
defer span.End()
3839

40+
ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
41+
defer cancel()
42+
3943
heartbeat := new(entities.Heartbeat)
4044
err := repository.db.WithContext(ctx).
4145
Where("user_id = ?", userID).
@@ -60,6 +64,9 @@ func (repository *gormHeartbeatRepository) Index(ctx context.Context, userID ent
6064
ctx, span := repository.tracer.Start(ctx)
6165
defer span.End()
6266

67+
ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
68+
defer cancel()
69+
6370
query := repository.db.WithContext(ctx).Where("user_id = ?", userID).Where("owner = ?", owner)
6471
if len(params.Query) > 0 {
6572
queryPattern := "%" + params.Query + "%"
@@ -80,6 +87,9 @@ func (repository *gormHeartbeatRepository) Store(ctx context.Context, heartbeat
8087
ctx, span := repository.tracer.Start(ctx)
8188
defer span.End()
8289

90+
ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
91+
defer cancel()
92+
8393
if err := repository.db.WithContext(ctx).Create(heartbeat).Error; err != nil {
8494
msg := fmt.Sprintf("cannot save heartbeat with ID [%s]", heartbeat.ID)
8595
return repository.tracer.WrapErrorSpan(span, stacktrace.Propagate(err, msg))

0 commit comments

Comments
 (0)