diff --git a/model/miner.go b/model/miner.go index 7749daf0..5e96ecde 100644 --- a/model/miner.go +++ b/model/miner.go @@ -7,11 +7,12 @@ import ( type StorageMiner struct { gorm.Model - Address util.DbAddr `gorm:"unique"` - Suspended bool - SuspendedReason string - Name string - Version string - Location string - Owner uint + Address util.DbAddr `gorm:"unique"` + Suspended bool + SuspendedReason string + Name string + Version string + Location string + Owner uint + ShuttlePreference []ShuttlePreference } diff --git a/model/shuttle.go b/model/shuttle.go index e37f55e5..97e51c73 100644 --- a/model/shuttle.go +++ b/model/shuttle.go @@ -8,12 +8,13 @@ import ( type Shuttle struct { gorm.Model - Handle string `gorm:"unique"` - Token string - Host string - PeerID string - LastConnection time.Time - Private bool - Open bool - Priority int + Handle string `gorm:"unique"` + Token string + Host string + PeerID string + LastConnection time.Time + Private bool + Open bool + Priority int + ShuttlePreference []ShuttlePreference } diff --git a/model/shuttle_preference.go b/model/shuttle_preference.go new file mode 100644 index 00000000..a2198800 --- /dev/null +++ b/model/shuttle_preference.go @@ -0,0 +1,10 @@ +package model + +import "gorm.io/gorm" + +type ShuttlePreference struct { + gorm.Model + StorageMinerID uint `gorm:"index:idx_pref"` + ShuttleID uint + Priority uint `gorm:"index:idx_pref, sort:asc"` +} diff --git a/shuttlepreferencemgr/shuttlepreferencemgr.go b/shuttlepreferencemgr/shuttlepreferencemgr.go new file mode 100644 index 00000000..049912a1 --- /dev/null +++ b/shuttlepreferencemgr/shuttlepreferencemgr.go @@ -0,0 +1,78 @@ +package shuttlepreferencemgr + +import ( + "context" + "sort" + + "github.com/application-research/estuary/model" + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +type ShuttlePreferenceManager struct { + DB *gorm.DB + MaxPriority int +} + +type StorageProviderShuttlePing struct { + StorageMinerID uint + pings []ShuttlePing +} + +type ShuttlePing struct { + ShuttleID uint + ping float64 +} + +func NewShuttlePreferenceManager(db *gorm.DB, maxPriority int) (*ShuttlePreferenceManager, error) { + spref := &ShuttlePreferenceManager{ + DB: db, + MaxPriority: maxPriority, + } + + return spref, nil +} + +// Get SPs and priority in reference to a given shuttle +func (spm *ShuttlePreferenceManager) Query(sourceShuttle string) ([]model.ShuttlePreference, error) { + var spl []model.ShuttlePreference + + if err := spm.DB.Find(&spl, "ShuttleID = ?", sourceShuttle).Error; err != nil { + return nil, err + } + + return spl, nil +} + +// Creates or update SP/Shuttle preference records associated with the provided SP ids +func (spm *ShuttlePreferenceManager) PostSP(ctx context.Context, pings StorageProviderShuttlePing) error { + + var toCreate []model.ShuttlePreference + sorted := sortShuttles(pings.pings) + + // Create entries for each SP/Shuttle Priority mapping + for pri, shuttle := range sorted { + // Only create records for this max priority # + if pri == spm.MaxPriority { + break + } + toCreate = append(toCreate, model.ShuttlePreference{ + StorageMinerID: pings.StorageMinerID, + ShuttleID: shuttle.ShuttleID, + Priority: uint(pri), + }) + } + + // Save entry - Upsert -> Update existing entries (if SP's shuttle priority changes) + spm.DB.Create(&toCreate).Clauses(clause.OnConflict{UpdateAll: true}) + return nil +} + +// Sorts a list of shuttles in increasing order of their ping +func sortShuttles(m []ShuttlePing) []ShuttlePing { + sort.Slice(m, func(i, j int) bool { + return m[i].ping < m[j].ping + }) + + return m +}