// Copyright 2019 Yunion // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. package quotas import ( "context" "database/sql" "strings" "sync" "yunion.io/x/jsonutils" "yunion.io/x/log" "yunion.io/x/pkg/errors" "yunion.io/x/onecloud/pkg/appsrv" "yunion.io/x/onecloud/pkg/cloudcommon/consts" "yunion.io/x/onecloud/pkg/mcclient" "yunion.io/x/onecloud/pkg/util/ctx" ) var ( usageCalculateWorker = appsrv.NewWorkerManager("usageCalculateWorker", 1, 1024, true) realTimeUsageCalculateWorker = appsrv.NewWorkerManager("realTimeUsageCalculateWorker", 1, 1024, true) usageDirtyMap = make(map[string]bool, 0) usageDirtyMapLock = &sync.Mutex{} ) func setDirty(key string) { usageDirtyMapLock.Lock() defer usageDirtyMapLock.Unlock() usageDirtyMap[key] = true } func clearDirty(key string) { usageDirtyMapLock.Lock() defer usageDirtyMapLock.Unlock() delete(usageDirtyMap, key) } func isDirty(key string) bool { usageDirtyMapLock.Lock() defer usageDirtyMapLock.Unlock() if _, ok := usageDirtyMap[key]; ok { return true } return false } type quotaTask struct { manager *SQuotaBaseManager keys IQuotaKeys usageChan chan IQuota key string } func (t *quotaTask) Run() { ctx := context.WithValue(ctx.CtxWithTime(), "task", t) usage := t.manager.newQuota() if !isDirty(t.key) { if t.usageChan != nil { t.manager.usageStore.GetQuota(ctx, t.keys, usage) t.usageChan <- usage } return } usage.SetKeys(t.keys) err := usage.FetchUsage(ctx) if err != nil { log.Debugf("usage.FetchUsage fail %s", err) if t.usageChan != nil { t.usageChan <- nil } return } t.manager.usageStore.SetQuota(ctx, nil, usage) clearDirty(t.key) if t.usageChan != nil { t.usageChan <- usage } } func (t *quotaTask) Dump() string { return "" } func (manager *SQuotaBaseManager) PostUsageJob(keys IQuotaKeys, usageChan chan IQuota, realTime bool) { if !consts.EnableQuotaCheck() { go func() { usageChan <- nil }() return } key := QuotaKeyString(keys) setDirty(key) var worker *appsrv.SWorkerManager if realTime { worker = realTimeUsageCalculateWorker } else { worker = usageCalculateWorker } task := quotaTask{ manager: manager, keys: keys, usageChan: usageChan, key: key, } worker.Run(&task, nil, nil) } func (manager *SQuotaBaseManager) CalculateQuotaUsages(ctx context.Context, userCred mcclient.TokenCredential, isStart bool) { if !consts.EnableQuotaCheck() { return } log.Infof("CalculateQuotaUsages") quota := manager.newQuota() keys := quota.GetKeys() keyFields := keys.Fields() q := manager.Query(keyFields...) keyList := make([]IQuotaKeys, 0) rows, err := q.Rows() if err != nil { if errors.Cause(err) != sql.ErrNoRows { log.Errorf("query quotas fail %s", err) } return } defer rows.Close() for rows.Next() { quota := manager.newQuota() err = q.Row2Struct(rows, quota) if err != nil { log.Errorf("Row2Struct fail %s", err) return } keyList = append(keyList, quota.GetKeys()) } var fields []string idNameMap, _ := manager.keyList2IdNameMap(ctx, keyList) log.Debugf("%s", jsonutils.Marshal(idNameMap)) for _, keys := range keyList { if idNameMap != nil { // no error, do check if len(fields) == 0 { fields = keys.Fields() } values := keys.Values() for i := range fields { if strings.HasSuffix(fields[i], "_id") && len(values[i]) > 0 && len(idNameMap[fields[i]][values[i]]) == 0 { log.Infof("%s=%s found not exists, delete quota with key %s", fields[i], values[i], jsonutils.Marshal(keys)) manager.DeleteAllQuotas(ctx, userCred, keys) manager.pendingStore.DeleteAllQuotas(ctx, userCred, keys) manager.usageStore.DeleteAllQuotas(ctx, userCred, keys) continue } } } manager.PostUsageJob(keys, nil, false) } } func (manager *SQuotaBaseManager) keyList2IdNameMap(ctx context.Context, keyList []IQuotaKeys) (map[string]map[string]string, error) { idMap := make(map[string]map[string]string) var fields []string for _, keys := range keyList { if len(fields) == 0 { fields = keys.Fields() } values := keys.Values() for i := range fields { if strings.HasSuffix(fields[i], "_id") && len(values[i]) > 0 { if _, ok := idMap[fields[i]]; !ok { idMap[fields[i]] = make(map[string]string) } if _, ok := idMap[fields[i]][values[i]]; !ok { idMap[fields[i]][values[i]] = "" } } } } ret, err := manager.GetIQuotaManager().FetchIdNames(ctx, idMap) if err != nil { return nil, err } else { return ret, nil } }