Adding Admin Client calls GetQuota & SetQuota (#59)

* Adding Admin Client calls `GetQuota` & `SetQuota`

This change set adds admin client calls to fetch and
mutate the OwnerRole quota[cpu,ram,disk].
This commit is contained in:
Robert Allen 2018-03-07 16:24:27 -08:00 committed by Renan DelValle
parent 66809c55f7
commit c0d2969976
3 changed files with 77 additions and 1 deletions

View file

@ -71,6 +71,8 @@ type Realis interface {
DrainHosts(hosts ...string) (*aurora.Response, *aurora.DrainHostsResult_, error)
EndMaintenance(hosts ...string) (*aurora.Response, *aurora.EndMaintenanceResult_, error)
MaintenanceStatus(hosts ...string) (*aurora.Response, *aurora.MaintenanceStatusResult_, error)
SetQuota(role string, cpu *float64, ram *int64, disk *int64) (*aurora.Response, error)
GetQuota(role string) (*aurora.Response, error)
}
type realisClient struct {
@ -944,3 +946,44 @@ func (r *realisClient) MaintenanceStatus(hosts ...string) (*aurora.Response, *au
return resp, result, nil
}
// SetQuota sets a quota aggregate for the given role
// TODO(zircote) Currently investigating an error that is returned from thrift calls that include resources for `NamedPort` and `NumGpu`
func (r *realisClient) SetQuota(role string, cpu *float64, ramMb *int64, diskMb *int64) (*aurora.Response, error) {
ram := aurora.NewResource()
ram.RamMb = ramMb
c := aurora.NewResource()
c.NumCpus = cpu
d := aurora.NewResource()
d.DiskMb = diskMb
quota := aurora.NewResourceAggregate()
quota.Resources = make(map[*aurora.Resource]bool)
quota.Resources[ram] = true
quota.Resources[c] = true
quota.Resources[d] = true
resp, retryErr := r.thriftCallWithRetries(func() (*aurora.Response, error) {
resp, retryErr := r.adminClient.SetQuota(role, quota)
if retryErr != nil {
return nil, errors.Wrap(retryErr, "Unable to set role quota")
}
return resp, nil
})
return resp, retryErr
}
// GetQuota returns the resource aggregate for the given role
func (r *realisClient) GetQuota(role string) (*aurora.Response, error) {
resp, retryErr := r.thriftCallWithRetries(func() (*aurora.Response, error) {
resp, retryErr :=r.adminClient.GetQuota(role)
if retryErr != nil {
return nil, errors.Wrap(retryErr, "Unable to get role quota")
}
return resp, nil
})
return resp, retryErr
}