github.com/letsencrypt/trillian@v1.1.2-0.20180615153820-ae375a99d36a/server/quota.go (about) 1 // Copyright 2017 Google Inc. All Rights Reserved. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package server 16 17 import ( 18 "flag" 19 "fmt" 20 "sync" 21 22 "github.com/golang/glog" 23 "github.com/google/trillian/quota" 24 ) 25 26 const ( 27 // QuotaNoop represents the noop quota implementation. 28 QuotaNoop = "noop" 29 ) 30 31 // NewQuotaManagerFunc is the signature of a function which can be registered 32 // to provide instances of a quota manager. 33 type NewQuotaManagerFunc func() (quota.Manager, error) 34 35 var ( 36 // QuotaSystem is a flag specifying which quota system is in use. 37 QuotaSystem = flag.String("quota_system", "mysql", fmt.Sprintf("Quota system to use. One of: %v", quotaSystems())) 38 39 qpMu sync.RWMutex 40 qpByName map[string]NewQuotaManagerFunc 41 ) 42 43 func init() { 44 if err := RegisterQuotaManager(QuotaNoop, func() (quota.Manager, error) { 45 return quota.Noop(), nil 46 }); err != nil { 47 glog.Fatalf("Failed to register %v: %v", QuotaNoop, err) 48 } 49 } 50 51 // RegisterQuotaManager registers the provided QuotaManager. 52 func RegisterQuotaManager(name string, qp NewQuotaManagerFunc) error { 53 qpMu.Lock() 54 defer qpMu.Unlock() 55 56 if qpByName == nil { 57 qpByName = make(map[string]NewQuotaManagerFunc) 58 } 59 60 _, exists := qpByName[name] 61 if exists { 62 return fmt.Errorf("quota provider %v already registered", name) 63 } 64 qpByName[name] = qp 65 return nil 66 } 67 68 // quotaSystems returns a slice of registered quota system names. 69 func quotaSystems() []string { 70 qpMu.RLock() 71 defer qpMu.RUnlock() 72 73 r := []string{} 74 for k := range qpByName { 75 r = append(r, k) 76 } 77 78 return r 79 } 80 81 // NewQuotaManagerFromFlags returns a quota.Manager implementation as speficied by flag. 82 func NewQuotaManagerFromFlags() (quota.Manager, error) { 83 return NewQuotaManager(*QuotaSystem) 84 } 85 86 // NewQuotaManager returns a quota.Manager implementation. 87 func NewQuotaManager(name string) (quota.Manager, error) { 88 qpMu.RLock() 89 defer qpMu.RUnlock() 90 91 f, exists := qpByName[name] 92 if !exists { 93 return nil, fmt.Errorf("unknown quota system: %v", name) 94 } 95 return f() 96 }