Skip to content

Commit 86a5d77

Browse files
johnstcnjohnstcn
committed
address comments
1 parent b7bb176 commit 86a5d77

File tree

5 files changed

+69
-27
lines changed

5 files changed

+69
-27
lines changed

‎coderd/database/dbmem/dbmem.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5958,6 +5958,9 @@ func (q *FakeQuerier) UpdateProvisionerDaemonLastSeenAt(_ context.Context, arg d
59585958
if q.provisionerDaemons[idx].ID != arg.ID {
59595959
continue
59605960
}
5961+
if q.provisionerDaemons[idx].LastSeenAt.Time.After(arg.LastSeenAt.Time) {
5962+
continue
5963+
}
59615964
q.provisionerDaemons[idx].LastSeenAt = arg.LastSeenAt
59625965
return nil
59635966
}

‎coderd/database/queries.sql.go

Lines changed: 7 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/provisionerdaemons.sql

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,4 +47,10 @@ WHERE
4747
RETURNING *;
4848

4949
-- name: UpdateProvisionerDaemonLastSeenAt :exec
50-
UPDATE provisioner_daemons SET last_seen_at = @last_seen_at WHERE id = @id;
50+
UPDATE provisioner_daemons
51+
SET
52+
last_seen_at = @last_seen_at
53+
WHERE
54+
id = @id
55+
AND
56+
last_seen_at <= @last_seen_at;

‎coderd/provisionerdserver/provisionerdserver.go

Lines changed: 29 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -208,8 +208,7 @@ func NewServer(
208208
HeartbeatFn: options.HeartbeatFn,
209209
}
210210

211-
go s.heartbeat()
212-
211+
go s.heartbeatLoop()
213212
return s, nil
214213
}
215214

@@ -222,44 +221,55 @@ func (s *server) timeNow() time.Time {
222221
return dbtime.Now()
223222
}
224223

225-
// heartbeat runs heartbeatOnce at the interval specified by HeartbeatInterval
224+
// heartbeatLoop runs heartbeatOnce at the interval specified by HeartbeatInterval
226225
// until the lifecycle context is canceled.
227-
func (s *server) heartbeat() {
226+
func (s *server) heartbeatLoop() {
228227
tick := time.NewTicker(time.Nanosecond)
229228
defer tick.Stop()
230229
for {
231230
select {
232231
case <-s.lifecycleCtx.Done():
232+
s.Logger.Debug(s.lifecycleCtx, "heartbeat loop canceled")
233233
return
234234
case <-tick.C:
235+
if s.lifecycleCtx.Err() != nil {
236+
return
237+
}
238+
start := s.timeNow()
235239
hbCtx, hbCancel := context.WithTimeout(s.lifecycleCtx, s.HeartbeatInterval)
236-
if err := s.heartbeatOnce(hbCtx); err != nil {
240+
if err := s.heartbeat(hbCtx); err != nil {
237241
if !xerrors.Is(err, context.DeadlineExceeded) && !xerrors.Is(err, context.Canceled) {
238242
s.Logger.Error(hbCtx, "heartbeat failed", slog.Error(err))
239243
}
240244
}
241245
hbCancel()
242-
tick.Reset(s.HeartbeatInterval)
246+
elapsed := s.timeNow().Sub(start)
247+
nextBeat := s.HeartbeatInterval - elapsed
248+
// avoid negative interval
249+
if nextBeat <= 0 {
250+
nextBeat = time.Nanosecond
251+
}
252+
tick.Reset(nextBeat)
243253
}
244254
}
245255
}
246256

247-
// heartbeatOnce updates the last seen at timestamp in the database.
257+
// heartbeat updates the last seen at timestamp in the database.
248258
// If HeartbeatFn is set, it will be called instead.
249-
func (s *server) heartbeatOnce(ctx context.Context) error {
250-
if s.HeartbeatFn != nil {
251-
return s.HeartbeatFn(ctx)
252-
}
253-
254-
if s.lifecycleCtx.Err() != nil {
259+
func (s *server) heartbeat(ctx context.Context) error {
260+
select {
261+
case <-ctx.Done():
255262
return nil
263+
default:
264+
if s.HeartbeatFn != nil {
265+
return s.HeartbeatFn(ctx)
266+
}
267+
//nolint:gocritic // This is specifically for updating the last seen at timestamp.
268+
return s.Database.UpdateProvisionerDaemonLastSeenAt(dbauthz.AsSystemRestricted(ctx), database.UpdateProvisionerDaemonLastSeenAtParams{
269+
ID: s.ID,
270+
LastSeenAt: sql.NullTime{Time: s.timeNow(), Valid: true},
271+
})
256272
}
257-
258-
//nolint:gocritic // This is specifically for updating the last seen at timestamp.
259-
return s.Database.UpdateProvisionerDaemonLastSeenAt(dbauthz.AsSystemRestricted(ctx), database.UpdateProvisionerDaemonLastSeenAtParams{
260-
ID: s.ID,
261-
LastSeenAt: sql.NullTime{Time: s.timeNow(), Valid: true},
262-
})
263273
}
264274

265275
// AcquireJob queries the database to lock a job.

‎coderd/provisionerdserver/provisionerdserver_test.go

Lines changed: 23 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -103,9 +103,15 @@ func TestHeartbeat(t *testing.T) {
103103
ctx, cancel := context.WithCancel(context.Background())
104104
t.Cleanup(cancel)
105105
heartbeatChan := make(chan struct{})
106-
heartbeatFn := func(context.Context) error {
107-
heartbeatChan <- struct{}{}
108-
return nil
106+
heartbeatFn := func(hbCtx context.Context) error {
107+
t.Logf("heartbeat")
108+
select {
109+
case <-hbCtx.Done():
110+
return hbCtx.Err()
111+
default:
112+
heartbeatChan <- struct{}{}
113+
return nil
114+
}
109115
}
110116
//nolint:dogsled // ｡:ﾟ૮ ˶ˆ ﻌ ˆ˶ ა ﾟ:｡
111117
_, _, _, _ = setup(t, false, &overrides{
@@ -114,10 +120,21 @@ func TestHeartbeat(t *testing.T) {
114120
heartbeatInterval: testutil.IntervalFast,
115121
})
116122

117-
<-heartbeatChan
123+
_, ok := <-heartbeatChan
124+
require.True(t, ok, "first heartbeat not received")
125+
_, ok = <-heartbeatChan
126+
require.True(t, ok, "second heartbeat not received")
118127
cancel()
128+
// Close the channel to ensure we don't receive any more heartbeats.
129+
// The test will fail if we do.
130+
defer func() {
131+
if r := recover(); r != nil {
132+
t.Fatalf("heartbeat received after cancel: %v", r)
133+
}
134+
}()
135+
119136
close(heartbeatChan)
120-
<-time.After(testutil.IntervalFast)
137+
<-time.After(testutil.IntervalMedium)
121138
}
122139

123140
func TestAcquireJob(t *testing.T) {
@@ -1766,7 +1783,7 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi
17661783
CreatedAt: dbtime.Now(),
17671784
Provisioners: []database.ProvisionerType{database.ProvisionerTypeEcho},
17681785
Tags: database.StringMap{},
1769-
LastSeenAt: sql.NullTime{Time: dbtime.Now(), Valid: true},
1786+
LastSeenAt: sql.NullTime{},
17701787
Version: "",
17711788
APIVersion: "1.0",
17721789
})

0 commit comments

Comments
 (0)