harmony: Fix last_contact race
This commit is contained in:
parent
adb0fc3259
commit
ca4e2d71b6
@ -3,7 +3,6 @@ package resources
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"os/exec"
|
||||
"regexp"
|
||||
"runtime"
|
||||
@ -49,30 +48,34 @@ func Register(db *harmonydb.DB, hostnameAndPort string) (*Reg, error) {
|
||||
}
|
||||
ctx := context.Background()
|
||||
{ // Learn our owner_id while updating harmony_machines
|
||||
var ownerID []int
|
||||
err := db.Select(ctx, &ownerID, `SELECT id FROM harmony_machines WHERE host_and_port=$1`, hostnameAndPort)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("could not read from harmony_machines: %w", err)
|
||||
}
|
||||
var ownerID int
|
||||
gpuram := uint64(lo.Sum(reg.GpuRam))
|
||||
if len(ownerID) == 0 {
|
||||
err = db.QueryRow(ctx, `INSERT INTO harmony_machines
|
||||
(host_and_port, cpu, ram, gpu, gpuram) VALUES
|
||||
($1,$2,$3,$4,$5) RETURNING id`,
|
||||
hostnameAndPort, reg.Cpu, reg.Ram, reg.Gpu, gpuram).Scan(®.Resources.MachineID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
} else {
|
||||
reg.MachineID = ownerID[0]
|
||||
_, err := db.Exec(ctx, `UPDATE harmony_machines SET
|
||||
cpu=$1, ram=$2, gpu=$3, gpuram=$4 WHERE id=$5`,
|
||||
reg.Cpu, reg.Ram, reg.Gpu, gpuram, reg.Resources.MachineID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Upsert query with last_contact update, fetch the machine ID
|
||||
// (note this isn't a simple insert .. on conflict because host_and_port isn't unique)
|
||||
err := db.QueryRow(ctx, `
|
||||
WITH upsert AS (
|
||||
UPDATE harmony_machines
|
||||
SET cpu = $2, ram = $3, gpu = $4, gpuram = $5, last_contact = CURRENT_TIMESTAMP
|
||||
WHERE host_and_port = $1
|
||||
RETURNING id
|
||||
),
|
||||
inserted AS (
|
||||
INSERT INTO harmony_machines (host_and_port, cpu, ram, gpu, gpuram, last_contact)
|
||||
SELECT $1, $2, $3, $4, $5, CURRENT_TIMESTAMP
|
||||
WHERE NOT EXISTS (SELECT id FROM upsert)
|
||||
RETURNING id
|
||||
)
|
||||
SELECT id FROM upsert
|
||||
UNION ALL
|
||||
SELECT id FROM inserted;
|
||||
`, hostnameAndPort, reg.Cpu, reg.Ram, reg.Gpu, gpuram).Scan(&ownerID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
reg.MachineID = ownerID
|
||||
|
||||
cleaned := CleanupMachines(context.Background(), db)
|
||||
logger.Infow("Cleaned up machines", "count", cleaned)
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user