feat(objnode): s3 api qos

Signed-off-by: tangdeyi <tangdeyi@oppo.com>
This commit is contained in:
tangdeyi 2023-08-01 11:01:21 +08:00 committed by tangdeyi
parent be64b60191
commit 14f0f67c73
36 changed files with 2426 additions and 64 deletions

View File

@ -1,3 +1,17 @@
// Copyright 2023 The CubeFS Authors.
//
// 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 master
import (
@ -1797,3 +1811,17 @@ func parseRequestToUpdateDecommissionDiskFactor(r *http.Request) (factor float64
}
return strconv.ParseFloat(value, 64)
}
func parseS3QosReq(r *http.Request, req *proto.S3QosRequest) (err error) {
var body []byte
if body, err = ioutil.ReadAll(r.Body); err != nil {
return
}
if err = json.Unmarshal(body, &req); err != nil {
return
}
log.LogInfo("parseS3QosReq success.")
return
}

View File

@ -5921,6 +5921,7 @@ func (m *Server) handleLcNodeTaskResponse(w http.ResponseWriter, r *http.Request
sendOkReply(w, r, newSuccessHTTPReply(fmt.Sprintf("%v", http.StatusOK)))
m.cluster.handleLcNodeTaskResponse(tr.OperatorAddr, tr)
}
func (m *Server) SetBucketLifecycle(w http.ResponseWriter, r *http.Request) {
var (
bytes []byte
@ -5942,6 +5943,7 @@ func (m *Server) SetBucketLifecycle(w http.ResponseWriter, r *http.Request) {
_ = m.cluster.SetBucketLifecycle(&req)
sendOkReply(w, r, newSuccessHTTPReply(fmt.Sprintf("PutBucketLifecycleConfiguration successful ")))
}
func (m *Server) GetBucketLifecycle(w http.ResponseWriter, r *http.Request) {
var (
err error
@ -5962,6 +5964,7 @@ func (m *Server) GetBucketLifecycle(w http.ResponseWriter, r *http.Request) {
}
sendOkReply(w, r, newSuccessHTTPReply(lcConf))
}
func (m *Server) DelBucketLifecycle(w http.ResponseWriter, r *http.Request) {
var (
err error
@ -5980,6 +5983,7 @@ func (m *Server) DelBucketLifecycle(w http.ResponseWriter, r *http.Request) {
log.LogWarn(msg)
sendOkReply(w, r, newSuccessHTTPReply(msg))
}
func (m *Server) lcnodeInfo(w http.ResponseWriter, r *http.Request) {
if err := r.ParseForm(); err != nil {
sendErrReply(w, r, &proto.HTTPReply{Code: proto.ErrCodeParamError, Msg: err.Error()})
@ -6007,3 +6011,182 @@ func (m *Server) lcnodeInfo(w http.ResponseWriter, r *http.Request) {
sendErrReply(w, r, &proto.HTTPReply{Code: proto.ErrCodeParamError, Msg: "invalid op"})
}
}
func (m *Server) S3QosSet(w http.ResponseWriter, r *http.Request) {
var (
param = &proto.S3QosRequest{}
err error
)
metric := exporter.NewTPCnt(apiToMetricsName(proto.S3QoSSet))
defer func() {
doStatAndMetric(proto.S3QoSSet, metric, err, nil)
}()
if err = parseS3QosReq(r, param); err != nil {
log.LogErrorf("[S3QosSet] parse fail err [%v]", err)
sendErrReply(w, r, &proto.HTTPReply{Code: proto.ErrCodeParamError, Msg: err.Error()})
return
}
if !isS3QosConfigValid(param) {
sendErrReply(w, r, &proto.HTTPReply{Code: proto.ErrCodeParamError, Msg: "s3 qos param err"})
return
}
// set s3 qos quota
if param.Quota != 0 {
if strings.ToLower(param.Uid) == proto.DefaultUid {
param.Uid = proto.DefaultUid
}
param.Api = strings.ToLower(param.Api)
metadata := new(RaftCmd)
metadata.Op = opSyncS3QosSet
key := param.Api + keySeparator + param.Uid + keySeparator + param.Type
metadata.K = S3QoSPrefix + key
metadata.V = []byte(strconv.FormatUint(param.Quota, 10))
// raft sync
if err = m.cluster.submit(metadata); err != nil {
sendErrReply(w, r, newErrHTTPReply(err))
return
}
// memory cache
m.cluster.S3ApiQosQuota.Store(metadata.K, param.Quota)
}
// set s3 node num
if param.Nodes != 0 {
metadata := new(RaftCmd)
metadata.Op = opSyncS3QosSet
key := proto.S3Nodes
metadata.K = S3QoSPrefix + key
metadata.V = []byte(strconv.FormatUint(param.Nodes, 10))
// raft sync
if err = m.cluster.submit(metadata); err != nil {
sendErrReply(w, r, newErrHTTPReply(err))
return
}
// memory cache
m.cluster.S3ApiQosQuota.Store(metadata.K, param.Nodes)
}
sendOkReply(w, r, newSuccessHTTPReply("success"))
}
func (m *Server) S3QosGet(w http.ResponseWriter, r *http.Request) {
var (
err error
)
metric := exporter.NewTPCnt(apiToMetricsName(proto.S3QoSGet))
defer func() {
doStatAndMetric(proto.S3QoSGet, metric, err, nil)
}()
apiLimitConf := make(map[string]*proto.UserLimitConf, 0)
s3QosResponse := proto.S3QoSResponse{
ApiLimitConf: apiLimitConf,
}
// memory cache
m.cluster.S3ApiQosQuota.Range(func(key, value interface{}) bool {
k := key.(string)
v := value.(uint64)
api, uid, limitType, nodeNumKey, err := parseS3QoSKey(k)
if err != nil {
log.LogErrorf("[S3QosGet] parseS3QoSKey err [%v]", err)
return true
}
if nodeNumKey != "" {
s3QosResponse.Nodes = v
return true
}
if _, ok := apiLimitConf[api]; !ok {
bandWidthQuota := make(map[string]uint64, 0)
qpsQuota := make(map[string]uint64, 0)
concurrentQuota := make(map[string]uint64, 0)
userLimitConf := &proto.UserLimitConf{
BandWidthQuota: bandWidthQuota,
QPSQuota: qpsQuota,
ConcurrentQuota: concurrentQuota,
}
apiLimitConf[api] = userLimitConf
}
switch limitType {
case proto.FlowLimit:
apiLimitConf[api].BandWidthQuota[uid] = v
case proto.QPSLimit:
apiLimitConf[api].QPSQuota[uid] = v
case proto.ConcurrentLimit:
apiLimitConf[api].ConcurrentQuota[uid] = v
}
return true
})
log.LogDebugf("[S3QosGet] s3qosInfoMap %+v", s3QosResponse)
sendOkReply(w, r, newSuccessHTTPReply(s3QosResponse))
}
func (m *Server) S3QosDelete(w http.ResponseWriter, r *http.Request) {
var (
param = &proto.S3QosRequest{}
err error
)
metric := exporter.NewTPCnt(apiToMetricsName(proto.S3QoSDelete))
defer func() {
doStatAndMetric(proto.S3QoSDelete, metric, err, nil)
}()
if err = parseS3QosReq(r, param); err != nil {
log.LogErrorf("[S3QosSet] parse fail err [%v]", err)
sendErrReply(w, r, &proto.HTTPReply{Code: proto.ErrCodeParamError, Msg: err.Error()})
return
}
if !isS3QosConfigValid(param) {
sendErrReply(w, r, &proto.HTTPReply{Code: proto.ErrCodeParamError, Msg: "s3 qos param err"})
return
}
if strings.ToLower(param.Uid) == proto.DefaultUid {
param.Uid = proto.DefaultUid
}
param.Api = strings.ToLower(param.Api)
metadata := new(RaftCmd)
metadata.Op = opSyncS3QosDelete
key := param.Api + keySeparator + param.Uid + keySeparator + param.Type
metadata.K = S3QoSPrefix + key
// raft sync
if err = m.cluster.submit(metadata); err != nil {
sendErrReply(w, r, newErrHTTPReply(err))
return
}
// memory cache
m.cluster.S3ApiQosQuota.Delete(metadata.K)
sendOkReply(w, r, newSuccessHTTPReply("success"))
}
func parseS3QoSKey(key string) (api, uid, limitType, nodes string, err error) {
s3qosInfo := strings.TrimPrefix(key, S3QoSPrefix)
strs := strings.Split(s3qosInfo, keySeparator)
if len(strs) == 3 {
return strs[0], strs[1], strs[2], "", nil
}
if len(strs) == 1 && strs[0] == proto.S3Nodes {
return "", "", "", strs[0], nil
}
return "", "", "", "", errors.New("unexpected key")
}
func isS3QosConfigValid(param *proto.S3QosRequest) bool {
if param.Type != proto.FlowLimit && param.Type != proto.QPSLimit && param.Type != proto.ConcurrentLimit {
return false
}
if proto.IsS3PutApi(param.Api) {
return false
}
return true
}

View File

@ -97,6 +97,7 @@ type Cluster struct {
lcMgr *lifecycleManager
snapshotMgr *snapshotDelManager
DecommissionDiskFactor float64
S3ApiQosQuota *sync.Map // (api,uid,limtType) -> limitQuota
}
type followerReadManager struct {
@ -336,6 +337,7 @@ func newCluster(name string, leaderInfo *LeaderInfo, fsm *MetadataFsm, partition
c.lcMgr.cluster = c
c.snapshotMgr = newSnapshotManager()
c.snapshotMgr.cluster = c
c.S3ApiQosQuota = new(sync.Map)
return
}

View File

@ -280,6 +280,9 @@ const (
opSyncSetQuota uint32 = 0x41
opSyncDeleteQuota uint32 = 0x42
opSyncMulitVersion uint32 = 0x53
opSyncS3QosSet uint32 = 0x60
opSyncS3QosDelete uint32 = 0x61
)
const (
@ -298,6 +301,7 @@ const (
domainAcronym = "zoneDomain"
apiLimiterAcronym = "al"
lcConfigurationAcronym = "lc"
S3QoS = "s3qos"
maxDataPartitionIDKey = keySeparator + "max_dp_id"
maxMetaPartitionIDKey = keySeparator + "max_mp_id"
maxCommonIDKey = keySeparator + "max_common_id"
@ -331,6 +335,7 @@ const (
quotaPrefix = keySeparator + "quota" + keySeparator
lcNodePrefix = keySeparator + lcNodeAcronym + keySeparator
lcConfPrefix = keySeparator + lcConfigurationAcronym + keySeparator
S3QoSPrefix = keySeparator + S3QoS + keySeparator
)
// selector enum

View File

@ -731,6 +731,16 @@ func (m *Server) registerAPIRoutes(router *mux.Router) {
Path(proto.QuotaListAll).
HandlerFunc(m.ListQuotaAll)
// S3 API QoS Manager
router.NewRoute().Methods(http.MethodPut, http.MethodPost).
Path(proto.S3QoSSet).
HandlerFunc(m.S3QosSet)
router.NewRoute().Methods(http.MethodGet, http.MethodPost).
Path(proto.S3QoSGet).
HandlerFunc(m.S3QosGet)
router.NewRoute().Methods(http.MethodDelete, http.MethodPost).
Path(proto.S3QoSDelete).
HandlerFunc(m.S3QosDelete)
}
func (m *Server) registerHandler(router *mux.Router, model string, schema *graphql.Schema) {

View File

@ -226,6 +226,12 @@ func (m *Server) loadMetadata() {
}
log.LogInfo("action[loadLcNodes] end")
syslog.Println("action[loadMetadata] end")
log.LogInfo("action[loadS3QoSInfo] begin")
if err = m.cluster.loadS3ApiQosInfo(); err != nil {
panic(err)
}
log.LogInfo("action[loadS3QoSInfo] end")
}
func (m *Server) clearMetadata() {

View File

@ -130,7 +130,7 @@ func (mf *MetadataFsm) Apply(command []byte, index uint64) (resp interface{}, er
switch cmd.Op {
case opSyncDeleteDataNode, opSyncDeleteMetaNode, opSyncDeleteVol, opSyncDeleteDataPartition, opSyncDeleteMetaPartition,
opSyncDeleteUserInfo, opSyncDeleteAKUser, opSyncDeleteVolUser, opSyncDeleteQuota, opSyncDeleteLcNode, opSyncDeleteLcConf:
opSyncDeleteUserInfo, opSyncDeleteAKUser, opSyncDeleteVolUser, opSyncDeleteQuota, opSyncDeleteLcNode, opSyncDeleteLcConf, opSyncS3QosDelete:
if err = mf.delKeyAndPutIndex(cmd.K, cmdMap); err != nil {
panic(err)
}
@ -143,6 +143,7 @@ func (mf *MetadataFsm) Apply(command []byte, index uint64) (resp interface{}, er
panic(err)
}
default:
// sync put data
if err = mf.store.BatchPut(cmdMap, true); err != nil {
panic(err)
}

View File

@ -1472,6 +1472,27 @@ func (c *Cluster) loadQuota() (err error) {
return
}
// load s3api qos info to memory cache
func (c *Cluster) loadS3ApiQosInfo() (err error) {
keyPrefix := S3QoSPrefix
result, err := c.fsm.store.SeekForPrefix([]byte(keyPrefix))
if err != nil {
err = fmt.Errorf("loadS3ApiQosInfo get failed, err [%v]", err)
return err
}
for key, value := range result {
s3qosQuota, err := strconv.ParseUint(string(value), 10, 64)
if err != nil {
return err
}
log.LogDebugf("loadS3ApiQosInfo key[%v] value[%v]", key, s3qosQuota)
c.S3ApiQosQuota.Store(key, s3qosQuota)
}
return
}
func (c *Cluster) addBadDataPartitionIdMap(dp *DataPartition) {
if !dp.IsDecommissionRunning() {
return

View File

@ -79,6 +79,13 @@ func (o *ObjectNode) createBucketHandler(w http.ResponseWriter, r *http.Request)
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(userInfo.UserID, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(userInfo.UserID, param.apiName)
// get LocationConstraint if any
contentLenStr := r.Header.Get(ContentLength)
if contentLen, errConv := strconv.Atoi(contentLenStr); errConv == nil && contentLen > 0 {
@ -161,13 +168,20 @@ func (o *ObjectNode) deleteBucketHandler(w http.ResponseWriter, r *http.Request)
GetRequestID(r), bucket, param.AccessKey(), err)
return
}
var vol *Volume
if vol, err = o.getVol(bucket); err != nil {
log.LogErrorf("deleteBucketHandler: load volume fail: requestID(%v) volume(%v) err(%v)",
GetRequestID(r), bucket, err)
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
if !vol.IsEmpty() {
errorCode = BucketNotEmpty
return
@ -206,13 +220,20 @@ func (o *ObjectNode) listBucketsHandler(w http.ResponseWriter, r *http.Request)
}()
param := ParseRequestParam(r)
userInfo, err := o.getUserInfoByAccessKeyV2(param.AccessKey())
if err != nil {
var userInfo *proto.UserInfo
if userInfo, err = o.getUserInfoByAccessKeyV2(param.accessKey); err != nil {
log.LogErrorf("listBucketsHandler: get user info fail: requestID(%v) accessKey(%v) err(%v)",
GetRequestID(r), param.AccessKey(), err)
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(userInfo.UserID, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(userInfo.UserID, param.apiName)
type bucket struct {
XMLName xml.Name `xml:"Bucket"`
CreationDate string `xml:"CreationDate"`
@ -262,17 +283,25 @@ func (o *ObjectNode) getBucketLocationHandler(w http.ResponseWriter, r *http.Req
o.errorResponse(w, r, err, errorCode)
}()
var vol *Volume
param := ParseRequestParam(r)
if param.Bucket() == "" {
errorCode = InvalidBucketName
return
}
if _, err = o.getVol(param.Bucket()); err != nil {
if vol, err = o.getVol(param.Bucket()); err != nil {
log.LogErrorf("getBucketLocationHandler: load volume fail: requestID(%v) volume(%v) err(%v)",
GetRequestID(r), param.Bucket(), err)
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
location := LocationResponse{Location: o.region}
response, err := MarshalXMLEntity(location)
if err != nil {
@ -308,6 +337,13 @@ func (o *ObjectNode) getBucketTaggingHandler(w http.ResponseWriter, r *http.Requ
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
var xattrInfo *proto.XAttrInfo
if xattrInfo, err = vol.GetXAttr(bucketRootPath, XAttrKeyOSSTagging); err != nil {
log.LogErrorf("getBucketTaggingHandler: Volume get XAttr fail: requestID(%v) err(%v)", GetRequestID(r), err)
@ -354,6 +390,13 @@ func (o *ObjectNode) putBucketTaggingHandler(w http.ResponseWriter, r *http.Requ
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
var body []byte
if body, err = ioutil.ReadAll(r.Body); err != nil {
log.LogErrorf("putBucketTaggingHandler: read request body data fail: requestID(%v) err(%v)",
@ -410,6 +453,13 @@ func (o *ObjectNode) deleteBucketTaggingHandler(w http.ResponseWriter, r *http.R
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
if err = vol.DeleteXAttr(bucketRootPath, XAttrKeyOSSTagging); err != nil {
log.LogErrorf("deleteBucketTaggingHandler: delete tagging xattr fail: requestID(%v) volume(%v) err(%v)",
GetRequestID(r), param.Bucket(), err)

View File

@ -68,6 +68,13 @@ func (o *ObjectNode) createMultipleUploadHandler(w http.ResponseWriter, r *http.
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
var userInfo *proto.UserInfo
if userInfo, err = o.getUserInfoByAccessKeyV2(param.AccessKey()); err != nil {
log.LogErrorf("createMultipleUploadHandler: get user info fail: requestID(%v) accessKey(%v) err(%v)",
@ -75,20 +82,14 @@ func (o *ObjectNode) createMultipleUploadHandler(w http.ResponseWriter, r *http.
return
}
// system metadata
// Get the requested content-type.
// In addition to being used to manage data types, it is used to distinguish
// whether the request is to create a directory.
// metadata
contentType := r.Header.Get(ContentType)
// Get request header : content-disposition
contentDisposition := r.Header.Get(ContentDisposition)
// Get request header : Cache-Control
cacheControl := r.Header.Get(CacheControl)
if len(cacheControl) > 0 && !ValidateCacheControl(cacheControl) {
errorCode = InvalidCacheArgument
return
}
// Get request header : Expires
expires := r.Header.Get(Expires)
if len(expires) > 0 && !ValidateCacheExpires(expires) {
errorCode = InvalidCacheArgument
@ -197,6 +198,17 @@ func (o *ObjectNode) uploadPartHandler(w http.ResponseWriter, r *http.Request) {
requestMD5 = hex.EncodeToString(decoded)
}
// Verify ContentLength
length := GetContentLength(r)
if length > SinglePutLimit {
errorCode = EntityTooLarge
return
}
if length < 0 {
errorCode = MissingContentLength
return
}
var vol *Volume
if vol, err = o.getVol(param.Bucket()); err != nil {
log.LogErrorf("uploadPartHandler: load volume fail: requestID(%v) err(%v)",
@ -216,8 +228,22 @@ func (o *ObjectNode) uploadPartHandler(w http.ResponseWriter, r *http.Request) {
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
// Flow Control
var reader io.Reader
if length > DefaultFlowLimitSize {
reader = rateLimit.GetReader(vol.owner, param.apiName, r.Body)
} else {
reader = r.Body
}
var fsFileInfo *FSFileInfo
if fsFileInfo, err = vol.WritePart(param.Object(), uploadId, uint16(partNumberInt), r.Body); err != nil {
if fsFileInfo, err = vol.WritePart(param.Object(), uploadId, uint16(partNumberInt), reader); err != nil {
log.LogErrorf("uploadPartHandler: write part fail: requestID(%v) volume(%v) path(%v) uploadId(%v) part(%v) err(%v)",
GetRequestID(r), vol.Name(), param.Object(), uploadId, partNumberInt, err)
err = handleWritePartErr(err)
@ -279,6 +305,13 @@ func (o *ObjectNode) uploadPartCopyHandler(w http.ResponseWriter, r *http.Reques
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
// step2: extract params from req
srcBucket, srcObject, _, err := extractSrcBucketKey(r)
if err != nil {
@ -337,9 +370,14 @@ func (o *ObjectNode) uploadPartCopyHandler(w http.ResponseWriter, r *http.Reques
writer.CloseWithError(err)
}()
// step5: upload part by copy
var fsFileInfo *FSFileInfo
fsFileInfo, err = vol.WritePart(param.Object(), uploadId, uint16(partNumberInt), reader)
// step5: upload part by copy and flow control
var rd io.Reader
if copyLength > DefaultFlowLimitSize {
rd = rateLimit.GetReader(vol.owner, param.apiName, reader)
} else {
rd = reader
}
fsFileInfo, err := vol.WritePart(param.Object(), uploadId, uint16(partNumberInt), rd)
if err != nil {
log.LogErrorf("uploadPartCopyHandler: write part fail: requestID(%v) volume(%v) path(%v) uploadId(%v) part(%v) err(%v)",
GetRequestID(r), vol.Name(), param.Object(), uploadId, partNumberInt, err)
@ -445,6 +483,13 @@ func (o *ObjectNode) listPartsHandler(w http.ResponseWriter, r *http.Request) {
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
fsParts, nextMarker, isTruncated, err := vol.ListParts(param.Object(), uploadId, maxPartsInt, partNoMarkerInt)
if err != nil {
log.LogErrorf("listPartsHandler: list parts fail, requestID(%v) uploadID(%v) maxParts(%v) partNoMarker(%v) err(%v)",
@ -598,6 +643,13 @@ func (o *ObjectNode) completeMultipartUploadHandler(w http.ResponseWriter, r *ht
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
// get uploaded part info in request
var requestBytes []byte
requestBytes, err = ioutil.ReadAll(r.Body)
@ -720,6 +772,13 @@ func (o *ObjectNode) abortMultipartUploadHandler(w http.ResponseWriter, r *http.
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
// Abort multipart upload
if err = vol.AbortMultipart(param.Object(), uploadId); err != nil {
log.LogErrorf("abortMultipartUploadHandler: abort multipart fail: requestID(%v) uploadID(%v) err(%v)",
@ -781,6 +840,13 @@ func (o *ObjectNode) listMultipartUploadsHandler(w http.ResponseWriter, r *http.
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
fsUploads, nextKeyMarker, nextUploadIdMarker, IsTruncated, prefixes, err := vol.ListMultipartUploads(prefix, delimiter, keyMarker, uploadIdMarker, maxUploadsInt)
if err != nil {
log.LogErrorf("listMultipartUploadsHandler: list multipart uploads fail: requestID(%v) err(%v)",

View File

@ -66,6 +66,14 @@ func (o *ObjectNode) getObjectHandler(w http.ResponseWriter, r *http.Request) {
GetRequestID(r), param.Bucket(), param.Object(), err)
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
// parse http range option
var (
revRange bool
@ -273,7 +281,15 @@ func (o *ObjectNode) getObjectHandler(w http.ResponseWriter, r *http.Request) {
size = rangeUpper - rangeLower + 1
w.WriteHeader(http.StatusPartialContent)
}
err = vol.readFile(fileInfo.Inode, fileSize, param.Object(), w, offset, size)
// Flow Control
var writer io.Writer
if size > DefaultFlowLimitSize {
writer = rateLimit.GetResponseWriter(vol.owner, param.apiName, w)
} else {
writer = w
}
err = vol.readFile(fileInfo.Inode, fileSize, param.Object(), writer, offset, size)
if err != nil {
log.LogErrorf("getObjectHandler: read file fail: requestID(%v) volume(%v) path(%v) offset(%v) size(%v) err(%v)",
GetRequestID(r), param.Bucket(), param.Object(), offset, size, err)
@ -373,6 +389,13 @@ func (o *ObjectNode) headObjectHandler(w http.ResponseWriter, r *http.Request) {
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
// get object meta
var fileInfo *FSFileInfo
fileInfo, _, err = vol.ObjectMeta(param.Object())
@ -623,6 +646,11 @@ func (o *ObjectNode) deleteObjectsHandler(w http.ResponseWriter, r *http.Request
objectKeys = append(objectKeys, object.Key)
log.LogWarnf("deleteObjectsHandler: delete path: requestID(%v) remote(%v) volume(%v) path(%v)",
GetRequestID(r), getRequestIP(r), vol.Name(), object.Key)
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, DELETE_OBJECT); err != nil {
return
}
if err = vol.DeletePath(object.Key); err != nil {
log.LogErrorf("deleteObjectsHandler: delete object failed: requestID(%v) volume(%v) path(%v) err(%v)",
GetRequestID(r), vol.Name(), object.Key, err)
@ -636,6 +664,7 @@ func (o *ObjectNode) deleteObjectsHandler(w http.ResponseWriter, r *http.Request
GetRequestID(r), vol.Name(), object.Key)
deletedObjects = append(deletedObjects, Deleted{Key: object.Key})
}
rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
}
deleteResult := DeleteResult{
@ -706,6 +735,13 @@ func (o *ObjectNode) copyObjectHandler(w http.ResponseWriter, r *http.Request) {
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
// client can reset these system metadata: Content-Type, Content-Disposition
contentType := r.Header.Get(ContentType)
contentDisposition := r.Header.Get(ContentDisposition)
@ -771,6 +807,14 @@ func (o *ObjectNode) copyObjectHandler(w http.ResponseWriter, r *http.Request) {
}
return
}
if fileInfo.Size > SinglePutLimit {
errorCode = EntityTooLarge
return
}
if fileInfo.Size > SinglePutLimit {
errorCode = EntityTooLarge
return
}
// get header
copyMatch := r.Header.Get(XAmzCopySourceIfMatch)
@ -891,6 +935,14 @@ func (o *ObjectNode) getBucketV1Handler(w http.ResponseWriter, r *http.Request)
GetRequestID(r), param.Bucket(), err)
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
// get options
marker := r.URL.Query().Get(ParamMarker)
prefix := r.URL.Query().Get(ParamPrefix)
@ -1019,6 +1071,13 @@ func (o *ObjectNode) getBucketV2Handler(w http.ResponseWriter, r *http.Request)
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
// get options
prefix := r.URL.Query().Get(ParamPrefix)
maxKeys := r.URL.Query().Get(ParamMaxKeys)
@ -1207,6 +1266,13 @@ func (o *ObjectNode) putObjectHandler(w http.ResponseWriter, r *http.Request) {
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
// Check 'x-amz-tagging' header
// Reference: https://docs.aws.amazon.com/AmazonS3/latest/API/API_PutObject.html#API_PutObject_RequestSyntax
var tagging *Tagging
@ -1238,6 +1304,17 @@ func (o *ObjectNode) putObjectHandler(w http.ResponseWriter, r *http.Request) {
return
}
// Verify ContentLength
length := GetContentLength(r)
if length > SinglePutLimit {
errorCode = EntityTooLarge
return
}
if length < 0 {
errorCode = MissingContentLength
return
}
// Get the requested content-type.
// In addition to being used to manage data types, it is used to distinguish
// whether the request is to create a directory.
@ -1274,7 +1351,15 @@ func (o *ObjectNode) putObjectHandler(w http.ResponseWriter, r *http.Request) {
ObjectLock: objetLock,
}
var startPut = time.Now()
if fsFileInfo, err = vol.PutObject(param.Object(), r.Body, opt); err != nil {
// Flow Control
var reader io.Reader
if length > DefaultFlowLimitSize {
reader = rateLimit.GetReader(vol.owner, param.apiName, r.Body)
} else {
reader = r.Body
}
if fsFileInfo, err = vol.PutObject(param.Object(), reader, opt); err != nil {
log.LogErrorf("putObjectHandler: put object fail: requestId(%v) volume(%v) path(%v) remote(%v) err(%v)",
GetRequestID(r), vol.Name(), param.Object(), getRequestIP(r), err)
err = handlePutObjectErr(err)
@ -1357,6 +1442,13 @@ func (o *ObjectNode) deleteObjectHandler(w http.ResponseWriter, r *http.Request)
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
// Audit deletion
log.LogInfof("Audit: delete object: requestID(%v) remote(%v) volume(%v) path(%v)",
GetRequestID(r), getRequestIP(r), vol.Name(), param.Object())
@ -1396,6 +1488,14 @@ func (o *ObjectNode) getObjectTaggingHandler(w http.ResponseWriter, r *http.Requ
GetRequestID(r), param.Bucket(), err)
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
var xattrInfo *proto.XAttrInfo
if xattrInfo, err = vol.GetXAttr(param.object, XAttrKeyOSSTagging); err != nil {
log.LogErrorf("getObjectTaggingHandler: get volume XAttr fail: requestID(%v) err(%v)", GetRequestID(r), err)
@ -1445,6 +1545,13 @@ func (o *ObjectNode) putObjectTaggingHandler(w http.ResponseWriter, r *http.Requ
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
var requestBody []byte
if requestBody, err = ioutil.ReadAll(r.Body); err != nil {
log.LogErrorf("putObjectTaggingHandler: read request body data fail: requestID(%v) err(%v)",
@ -1504,6 +1611,14 @@ func (o *ObjectNode) deleteObjectTaggingHandler(w http.ResponseWriter, r *http.R
GetRequestID(r), err)
return
}
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
if err = vol.DeleteXAttr(param.object, XAttrKeyOSSTagging); err != nil {
log.LogErrorf("deleteObjectTaggingHandler: volume delete tagging fail: requestID(%v) volume(%v) object(%v) err(%v)",
GetRequestID(r), param.Bucket(), param.Object(), err)
@ -1527,16 +1642,24 @@ func (o *ObjectNode) putObjectXAttrHandler(w http.ResponseWriter, r *http.Reques
errorCode = InvalidBucketName
return
}
if len(param.Object()) == 0 {
errorCode = InvalidKey
return
}
var vol *Volume
if vol, err = o.getVol(param.bucket); err != nil {
log.LogErrorf("pubObjectXAttrHandler: load volume fail: requestID(%v) err(%v)",
GetRequestID(r), err)
return
}
if len(param.Object()) == 0 {
errorCode = InvalidKey
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
var requestBody []byte
if requestBody, err = ioutil.ReadAll(r.Body); err != nil {
errorCode = &ErrorCode{
@ -1585,16 +1708,24 @@ func (o *ObjectNode) getObjectXAttrHandler(w http.ResponseWriter, r *http.Reques
errorCode = InvalidBucketName
return
}
if len(param.Object()) == 0 {
errorCode = InvalidKey
return
}
var vol *Volume
if vol, err = o.getVol(param.Bucket()); err != nil {
log.LogErrorf("getObjectXAttrHandler: load volume fail: requestID(%v) err(%v)",
GetRequestID(r), err)
return
}
if len(param.Object()) == 0 {
errorCode = InvalidKey
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
var xattrKey string
if xattrKey = param.GetVar(ParamKey); len(xattrKey) == 0 {
errorCode = InvalidArgument
@ -1640,16 +1771,24 @@ func (o *ObjectNode) deleteObjectXAttrHandler(w http.ResponseWriter, r *http.Req
errorCode = InvalidBucketName
return
}
if len(param.Object()) == 0 {
errorCode = InvalidKey
return
}
var vol *Volume
if vol, err = o.getVol(param.Bucket()); err != nil {
log.LogErrorf("deleteObjectXAttrHandler: load volume fail: requestID(%v) err(%v)",
GetRequestID(r), err)
return
}
if len(param.Object()) == 0 {
errorCode = InvalidKey
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
var xattrKey string
if xattrKey = param.GetVar(ParamKey); len(xattrKey) == 0 {
errorCode = InvalidArgument
@ -1681,16 +1820,23 @@ func (o *ObjectNode) listObjectXAttrs(w http.ResponseWriter, r *http.Request) {
errorCode = InvalidBucketName
return
}
if len(param.Object()) == 0 {
errorCode = InvalidKey
return
}
var vol *Volume
if vol, err = o.getVol(param.bucket); err != nil {
log.LogErrorf("listObjectXAttrs: load volume fail: requestID(%v) err(%v)",
GetRequestID(r), err)
return
}
if len(param.Object()) == 0 {
errorCode = InvalidKey
// QPS and Concurrency Limit
rateLimit := o.AcquireRateLimiter()
if err = rateLimit.AcquireLimitResource(vol.owner, param.apiName); err != nil {
return
}
defer rateLimit.ReleaseLimitResource(vol.owner, param.apiName)
var keys []string
if keys, err = vol.ListXAttrs(param.object); err != nil {
@ -1805,3 +1951,14 @@ func parsePartInfo(partNumber uint64, fileSize uint64) (uint64, uint64, uint64,
}
return partSize, partCount, rangeLower, rangeUpper
}
func GetContentLength(r *http.Request) int64 {
dcl := r.Header.Get(HeaderNameXAmzDecodedContentLength)
if dcl != "" {
length, err := strconv.ParseInt(dcl, 10, 64)
if err == nil {
return length
}
}
return r.ContentLength
}

View File

@ -84,6 +84,8 @@ const (
XAmzSecurityToken = "X-Amz-Security-Token"
XAmzObjectLockMode = "X-Amz-Object-Lock-Mode"
XAmzObjectLockRetainUntilDate = "X-Amz-Object-Lock-Retain-Until-Date"
HeaderNameXAmzDecodedContentLength = "x-amz-decoded-content-length"
)
const (
@ -127,9 +129,10 @@ const (
)
const (
MaxKeys = 1000
MaxParts = 1000
MaxUploads = 1000
MaxKeys = 1000
MaxParts = 1000
MaxUploads = 1000
SinglePutLimit = 5 * 1 << 30 // 5G
)
const (

View File

@ -366,16 +366,18 @@ func (v *Volume) IsEmpty() bool {
func (v *Volume) GetXAttr(path string, key string) (info *proto.XAttrInfo, err error) {
var inode uint64
var notUseCache bool
if objMetaCache != nil {
var retry = 0
for {
if _, inode, _, _, err = v.recursiveLookupTarget(path); err != nil {
if _, inode, _, _, err = v.recursiveLookupTarget(path, notUseCache); err != nil {
return v.getXAttr(path, key)
}
_, err = v.mw.InodeGet_ll(inode)
if err == syscall.ENOENT && retry < MaxRetry {
notUseCache = true
retry++
continue
}
@ -465,15 +467,17 @@ func (v *Volume) listXAttrs(path string) (keys []string, err error) {
func (v *Volume) ListXAttrs(path string) (keys []string, err error) {
var inode uint64
var notUseCache bool
if objMetaCache != nil {
var retry = 0
for {
if _, inode, _, _, err = v.recursiveLookupTarget(path); err != nil {
if _, inode, _, _, err = v.recursiveLookupTarget(path, notUseCache); err != nil {
return v.listXAttrs(path)
}
_, err = v.mw.InodeGet_ll(inode)
if err == syscall.ENOENT && retry < MaxRetry {
notUseCache = true
retry++
continue
}
@ -873,7 +877,7 @@ func (v *Volume) DeletePath(path string) (err error) {
var ino uint64
var name string
var mode os.FileMode
parent, ino, name, mode, err = v.recursiveLookupTarget(path)
parent, ino, name, mode, err = v.recursiveLookupTarget(path, false)
if err != nil {
// An unexpected error occurred
return
@ -912,7 +916,6 @@ func (v *Volume) DeletePath(path string) (err error) {
return
}
// Evict inode
if err = v.ec.EvictStream(ino); err != nil {
log.LogWarnf("DeletePath EvictStream: path(%v) inode(%v)", path, ino)
}
@ -922,6 +925,7 @@ func (v *Volume) DeletePath(path string) (err error) {
deleteAttrCache(parent, v.name)
log.LogWarnf("DeletePath: evict: volume(%v) path(%v) inode(%v)", v.name, path, ino)
// Evict inode
if err = v.mw.Evict(ino); err != nil {
log.LogWarnf("DeletePath Evict: path(%v) inode(%v)", path, ino)
}
@ -1634,7 +1638,7 @@ func (v *Volume) ReadFile(path string, writer io.Writer, offset, size uint64) er
var ino uint64
var mode os.FileMode
if _, ino, _, mode, err = v.recursiveLookupTarget(path); err != nil {
if _, ino, _, mode, err = v.recursiveLookupTarget(path, false); err != nil {
return err
}
@ -1656,15 +1660,16 @@ func (v *Volume) ObjectMeta(path string) (info *FSFileInfo, xattr *proto.XAttrIn
var mode os.FileMode
var inoInfo *proto.InodeInfo
var retry = 0
var notUseCache bool
for {
if _, inode, _, mode, err = v.recursiveLookupTarget(path); err != nil {
if _, inode, _, mode, err = v.recursiveLookupTarget(path, notUseCache); err != nil {
log.LogErrorf("ObjectMeta: recursive look up path fail: volume(%v) path(%v) err(%v)",
v.name, path, err)
return
}
inoInfo, err = v.mw.InodeGet_ll(inode)
if err == syscall.ENOENT && retry < MaxRetry {
notUseCache = true
retry++
continue
}
@ -1687,19 +1692,18 @@ func (v *Volume) ObjectMeta(path string) (info *FSFileInfo, xattr *proto.XAttrIn
if objMetaCache != nil {
attrItem, needRefresh := objMetaCache.GetAttr(v.name, inode)
if attrItem == nil || needRefresh {
log.LogDebugf("ObjectMeta: get attr in cache failed: volume(%v) inode(%v) attrItem(%v), needRefresh(%v)",
log.LogDebugf("ObjectMeta: get attr in cache miss: volume(%v) inode(%v) attrItem(%v), needRefresh(%v)",
v.name, inode, attrItem, needRefresh)
xattr, err = v.mw.XAttrGetAll_ll(inode)
if err != nil {
log.LogErrorf("ObjectMeta: XAttrGetAll_ll fail, volume(%v) inode(%v) path(%v) err(%v)",
v.name, inode, path, err)
return
} else {
attrItem = &AttrItem{
XAttrInfo: *xattr,
}
objMetaCache.PutAttr(v.name, attrItem)
}
attrItem = &AttrItem{
XAttrInfo: *xattr,
}
objMetaCache.PutAttr(v.name, attrItem)
} else {
xattr = &proto.XAttrInfo{
XAttrs: attrItem.XAttrs,
@ -1795,7 +1799,7 @@ func (v *Volume) Close() error {
// ENOENT:
// 0x2 ENOENT No such file or directory. A component of a specified
// pathname did not exist, or the pathname was an empty string.
func (v *Volume) recursiveLookupTarget(path string) (parent uint64, ino uint64, name string, mode os.FileMode, err error) {
func (v *Volume) recursiveLookupTarget(path string, notUseCache bool) (parent uint64, ino uint64, name string, mode os.FileMode, err error) {
parent = rootIno
var pathIterator = NewPathIterator(path)
if !pathIterator.HasNext() {
@ -1805,7 +1809,7 @@ func (v *Volume) recursiveLookupTarget(path string) (parent uint64, ino uint64,
cacheUsed := false
if objMetaCache != nil {
if objMetaCache != nil && !notUseCache {
for pathIterator.HasNext() {
var pathItem = pathIterator.Next()
var curIno uint64
@ -2577,7 +2581,7 @@ func (v *Volume) CopyFile(sv *Volume, sourcePath, targetPath, metaDirective stri
sInodeInfo *proto.InodeInfo
)
if _, sInode, sName, sMode, err = sv.recursiveLookupTarget(sourcePath); err != nil {
if _, sInode, sName, sMode, err = sv.recursiveLookupTarget(sourcePath, false); err != nil {
log.LogErrorf("CopyFile: look up source path fail, source path(%v) err(%v)", sourcePath, err)
return
}
@ -2670,21 +2674,20 @@ func (v *Volume) CopyFile(sv *Volume, sourcePath, targetPath, metaDirective stri
// operation at target object
var (
tMode os.FileMode
tInode uint64
oldtInode uint64
tInodeInfo *proto.InodeInfo
tParentId uint64
pathItems []PathItem
tLastName string
)
if _, oldtInode, _, tMode, err = v.recursiveLookupTarget(targetPath); err != nil && err != syscall.ENOENT {
if _, oldtInode, _, tMode, err = v.recursiveLookupTarget(targetPath, false); err != nil && err != syscall.ENOENT {
log.LogErrorf("CopyFile: look up target path failed, target path(%v), err(%v)", targetPath, err)
return
}
// if target file existed, check target file mode is whether same with source file
if err != syscall.ENOENT && tMode.IsDir() != sMode.IsDir() {
log.LogErrorf("CopyFile: target path existed and target path mode not same with source path, "+
"target path(%v), target inode(%v), source path(%v), source inode(%v)", targetPath, tInode, sourcePath, sInode)
"target path(%v), target inode(%v), source path(%v), source inode(%v)", targetPath, oldtInode, sourcePath, sInode)
return nil, syscall.EINVAL
}
// if source file mode is directory, return OK, and need't create target directory
@ -2951,17 +2954,8 @@ func (v *Volume) CopyFile(sv *Volume, sourcePath, targetPath, metaDirective stri
}
// force updating dentry and attrs in cache
if objMetaCache != nil {
dentry := &DentryItem{
Dentry: metanode.Dentry{
ParentId: tParentId,
Name: tLastName,
Inode: tInodeInfo.Inode,
Type: DefaultFileMode,
},
}
objMetaCache.PutDentry(v.name, dentry)
}
updateDentryCache(tParentId, tInodeInfo.Inode, DefaultFileMode, tLastName, v.name)
putAttrCache(targetAttr, v.name)
return
}

325
objectnode/ratelimit.go Normal file
View File

@ -0,0 +1,325 @@
// Copyright 2023 The CubeFS Authors.
//
// 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 objectnode
import (
"encoding/json"
"io"
"strconv"
"strings"
"sync"
"time"
"github.com/cubefs/cubefs/proto"
"github.com/cubefs/cubefs/util/concurrent"
"github.com/cubefs/cubefs/util/flowctrl"
"github.com/cubefs/cubefs/util/ratelimit"
)
const (
PUT = "put"
List = "list"
DefaultFlowLimitSize = 4 * 1024 // 4kB
)
var putApi = map[string]string{
PUT_OBJECT: PUT,
COPY_OBJECT: PUT,
POST_OBJECT: PUT,
UPLOAD_PART: PUT,
UPLOAD_PART_COPY: PUT,
}
var listApi = map[string]string{
List_BUCKETS: List,
LIST_OBJECTS: List,
LIST_OBJECTS_V2: List,
LIST_PARTS: List,
LIST_MULTIPART_UPLOADS: List,
}
type RateLimiter interface {
AcquireLimitResource(uid string, api string) error
ReleaseLimitResource(uid string, api string)
GetResponseWriter(uid string, api string, w io.Writer) io.Writer
GetReader(uid string, api string, r io.Reader) io.Reader
}
type RateLimit struct {
S3ApiRateLimitMgr map[string]UserRateManager // api -> UserRateMgr
ApiLimitConf map[string]*proto.UserLimitConf // api -> UserLimitConf
putApi map[string]string
limitMutex sync.RWMutex
}
func NewRateLimit(apiLimitConf map[string]*proto.UserLimitConf) RateLimiter {
if len(apiLimitConf) == 0 {
return &NullRateLimit{}
}
s3ApiRateLimitMgr := make(map[string]UserRateManager, len(apiLimitConf))
for api, userLimitConf := range apiLimitConf {
userRateManager := NewUserRateMgr(userLimitConf)
s3ApiRateLimitMgr[api] = userRateManager
}
rateLimit := &RateLimit{
S3ApiRateLimitMgr: s3ApiRateLimitMgr,
ApiLimitConf: apiLimitConf,
putApi: putApi,
}
return rateLimit
}
func (r *RateLimit) AcquireLimitResource(uid string, api string) error {
api = strings.ToLower(api)
if putTotal, isPutApi := r.putApi[api]; isPutApi {
api = putTotal
}
userRateMgr, ok := r.S3ApiRateLimitMgr[api]
if !ok {
return nil
}
// QPS
if allowed, _ := userRateMgr.QPSLimitAllowed(uid); !allowed {
return TooManyRequests
}
// Concurrency
if err := userRateMgr.ConcurrentLimitAcquire(uid); err != nil {
return TooManyRequests
}
return nil
}
func (r *RateLimit) ReleaseLimitResource(uid string, api string) {
api = strings.ToLower(api)
if putTotal, isPutApi := r.putApi[api]; isPutApi {
api = putTotal
}
userRateMgr, ok := r.S3ApiRateLimitMgr[api]
if !ok {
return
}
userRateMgr.ConcurrentLimitRelease(uid)
}
func (r *RateLimit) GetResponseWriter(uid string, api string, w io.Writer) io.Writer {
api = strings.ToLower(api)
userRateMgr, ok := r.S3ApiRateLimitMgr[api]
if !ok {
return w
}
return userRateMgr.GetResponseWriter(uid, w)
}
func (r *RateLimit) GetReader(uid string, api string, reader io.Reader) io.Reader {
api = strings.ToLower(api)
if putTotal, isPutApi := r.putApi[api]; isPutApi {
api = putTotal
}
userRateMgr, ok := r.S3ApiRateLimitMgr[api]
if !ok {
return reader
}
return userRateMgr.GetReader(uid, reader)
}
// No RateLimit
type NullRateLimit struct {
}
func (n *NullRateLimit) AcquireLimitResource(uid string, api string) error {
return nil
}
func (n *NullRateLimit) ReleaseLimitResource(uid string, api string) {
return
}
func (n *NullRateLimit) GetResponseWriter(uid string, api string, w io.Writer) io.Writer {
return w
}
func (n *NullRateLimit) GetReader(uid string, api string, r io.Reader) io.Reader {
return r
}
type UserRateManager interface {
QPSLimitAllowed(uid string) (bool, time.Duration)
ConcurrentLimitAcquire(uid string) error
ConcurrentLimitRelease(uid string)
GetResponseWriter(uid string, w io.Writer) io.Writer
GetReader(uid string, r io.Reader) io.Reader
}
// UserRateMgr specific api user rate Manager
type UserRateMgr struct {
BandWidthLimit *flowctrl.KeyFlowCtrl
QPSLimit *ratelimit.KeyRateLimit
ConcurrentLimit *concurrent.KeyConcurrentLimit
UserLimitConf *proto.UserLimitConf
}
func NewUserRateMgr(conf *proto.UserLimitConf) UserRateManager {
bandWidthLimit := flowctrl.NewKeyFlowCtrl()
qpsLimit := ratelimit.NewKeyRateLimit()
concurrentLimit := concurrent.NewLimit()
userRateMgr := &UserRateMgr{
BandWidthLimit: bandWidthLimit,
QPSLimit: qpsLimit,
ConcurrentLimit: concurrentLimit,
UserLimitConf: conf,
}
return userRateMgr
}
func (r *UserRateMgr) QPSLimitAllowed(uid string) (bool, time.Duration) {
defaultQPSLimit := r.UserLimitConf.QPSQuota[proto.DefaultUid]
usrQPSLimit := r.UserLimitConf.QPSQuota[uid]
qpsQuota := getUserLimitQuota(defaultQPSLimit, usrQPSLimit)
if qpsQuota == 0 {
return true, 0
}
qpsLimit := r.QPSLimit.Acquire(uid, qpsQuota)
defer r.QPSLimit.Release(uid)
return !qpsLimit.Limit(), 0
}
func (r *UserRateMgr) ConcurrentLimitAcquire(uid string) error {
defaultConcurrentLimit := r.UserLimitConf.ConcurrentQuota[proto.DefaultUid]
usrConcurrentLimit := r.UserLimitConf.ConcurrentQuota[uid]
concurrentQuota := getUserLimitQuota(defaultConcurrentLimit, usrConcurrentLimit)
if concurrentQuota == 0 {
return nil
}
return r.ConcurrentLimit.Acquire(uid, int64(concurrentQuota))
}
func (r *UserRateMgr) ConcurrentLimitRelease(uid string) {
r.ConcurrentLimit.Release(uid)
}
func (r *UserRateMgr) GetResponseWriter(uid string, w io.Writer) io.Writer {
defaultBandWidthLimit := r.UserLimitConf.BandWidthQuota[proto.DefaultUid]
usrBandWidthLimit := r.UserLimitConf.BandWidthQuota[uid]
bandWidthQuota := getUserLimitQuota(defaultBandWidthLimit, usrBandWidthLimit)
if bandWidthQuota == 0 {
return w
}
rate, _ := convertUint64ToInt(bandWidthQuota)
flowCtrl := r.BandWidthLimit.Acquire(uid, rate)
defer r.BandWidthLimit.Release(uid)
w = flowctrl.NewRateWriterWithCtrl(w, flowCtrl)
return w
}
func (r *UserRateMgr) GetReader(uid string, reader io.Reader) io.Reader {
defaultBandWidthLimit := r.UserLimitConf.BandWidthQuota[proto.DefaultUid]
usrBandWidthLimit := r.UserLimitConf.BandWidthQuota[uid]
bandWidthQuota := getUserLimitQuota(defaultBandWidthLimit, usrBandWidthLimit)
if bandWidthQuota == 0 {
return reader
}
rate, _ := convertUint64ToInt(bandWidthQuota)
flowCtrl := r.BandWidthLimit.Acquire(uid, rate)
defer r.BandWidthLimit.Release(uid)
reader = flowctrl.NewRateReaderWithCtrl(reader, flowCtrl)
return reader
}
// priority: usrLimit > defaultLimit
func getUserLimitQuota(defaultLimit, usrLimit uint64) uint64 {
if usrLimit != 0 {
return usrLimit
}
return defaultLimit
}
func (o *ObjectNode) Reload(data []byte) error {
s3QosResponse := proto.S3QoSResponse{}
if err := json.Unmarshal(data, &s3QosResponse); err != nil {
return err
}
apiLimitConf := s3QosResponse.ApiLimitConf
s3NodeNum := s3QosResponse.Nodes
for _, userLimitConf := range apiLimitConf {
for uid, bandWidthQuota := range userLimitConf.BandWidthQuota {
quota := bandWidthQuota / s3NodeNum
if quota == 0 && bandWidthQuota != 0 {
quota += 1
}
userLimitConf.BandWidthQuota[uid] = quota
}
for uid, qpsQuota := range userLimitConf.QPSQuota {
quota := qpsQuota / s3NodeNum
if quota == 0 && qpsQuota != 0 {
quota += 1
}
userLimitConf.QPSQuota[uid] = quota
}
for uid, concurrentQuota := range userLimitConf.ConcurrentQuota {
quota := concurrentQuota / s3NodeNum
if quota == 0 && concurrentQuota != 0 {
quota += 1
}
userLimitConf.ConcurrentQuota[uid] = quota
}
}
o.limitMutex.Lock()
o.rateLimit = NewRateLimit(apiLimitConf)
o.limitMutex.Unlock()
return nil
}
func (o *ObjectNode) requestRemote() (data []byte, err error) {
data, err = o.mc.AdminAPI().GetS3QoSInfo()
if err != nil {
return nil, err
}
return data, nil
}
func (o *ObjectNode) AcquireRateLimiter() RateLimiter {
o.limitMutex.RLock()
rateLimit := o.rateLimit
o.limitMutex.RUnlock()
return rateLimit
}
func convertUint64ToInt(num uint64) (int, error) {
str := strconv.FormatUint(num, 10)
parsed, err := strconv.ParseInt(str, 10, 0)
if err != nil {
return 0, err
}
return int(parsed), nil
}

View File

@ -87,6 +87,7 @@ var (
NoSuchObjectLockConfiguration = &ErrorCode{"NoSuchObjectLockConfiguration", "The specified object does not have a ObjectLock configuration", http.StatusNotFound}
NoContentMd5HeaderErr = &ErrorCode{"NoContentMd5Header", "Content-MD5 HTTP header is required for Upload Object/Part requests with Object Lock parameters", http.StatusBadRequest}
ObjectLockConfigurationNotFound = &ErrorCode{"ObjectLockConfigurationNotFoundError", "Object Lock configuration does not exist for this bucket", http.StatusNotFound}
TooManyRequests = &ErrorCode{"TooManyRequests", "too many requests, please retry later", http.StatusTooManyRequests}
)
type ErrorCode struct {

View File

@ -32,6 +32,7 @@ import (
"github.com/cubefs/cubefs/util/config"
"github.com/cubefs/cubefs/util/exporter"
"github.com/cubefs/cubefs/util/log"
"github.com/cubefs/cubefs/util/reloadconf"
"github.com/gorilla/mux"
)
@ -121,10 +122,12 @@ const (
// Default of configuration value
const (
defaultListen = "80"
defaultCacheRefreshInterval = 10 * 60
defaultMaxDentryCacheNum = 1000000
defaultMaxInodeAttrCacheNum = 1000000
defaultListen = "80"
defaultCacheRefreshInterval = 10 * 60
defaultMaxDentryCacheNum = 1000000
defaultMaxInodeAttrCacheNum = 1000000
defaultS3QoSReloadIntervalMs = 300 * 1000
defaultS3QoSConfName = "s3qosInfo.conf"
// ebs
MaxSizePutOnce = int64(1) << 23
)
@ -157,7 +160,9 @@ type ObjectNode struct {
disabledActions proto.Actions // disabled actions
stsNotAllowedActions proto.Actions // actions that are not accessible to STS users
control common.Control
control common.Control
rateLimit RateLimiter
limitMutex sync.RWMutex
}
func (o *ObjectNode) Start(cfg *config.Config) (err error) {
@ -303,6 +308,17 @@ func handleStart(s common.Server, cfg *config.Config) (err error) {
readThreads = rt
}
}
// s3 api qos info
reloadConf := &reloadconf.ReloadConf{
ConfName: defaultS3QoSConfName,
ReloadMs: defaultS3QoSReloadIntervalMs,
RequestRemote: o.requestRemote,
}
err = reloadconf.StartReload(reloadConf, o.Reload)
if err != nil {
log.LogWarnf("handleStart: GetS3QoSInfo err(%v)", err)
}
// start rest api
if err = o.startMuxRestAPI(); err != nil {

View File

@ -224,6 +224,11 @@ const (
QuotaGet = "/quota/get"
// QuotaBatchModifyPath = "/quota/batchModifyPath"
QuotaListAll = "/quota/listAll"
// s3 qos api
S3QoSSet = "/s3/qos/set"
S3QoSGet = "/s3/qos/get"
S3QoSDelete = "/s3/qos/delete"
)
var GApiInfo map[string]string = map[string]string{

53
proto/s3qos.go Normal file
View File

@ -0,0 +1,53 @@
// Copyright 2023 The CubeFS Authors.
//
// 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 proto
import "strings"
const (
FlowLimit = "f"
QPSLimit = "q"
ConcurrentLimit = "c"
S3Nodes = "s3nodes"
DefaultUid = "default"
)
type UserLimitConf struct {
BandWidthQuota map[string]uint64 `json:"band_width_quota"` // uid --> BytesPS
QPSQuota map[string]uint64 `json:"qps_quota"` // uid --> QPS
ConcurrentQuota map[string]uint64 `json:"concurrent_quota"` // uid --> concurrency
}
type S3QosRequest struct {
Uid string `json:"uid"`
Api string `json:"api"`
Type string `json:"type"`
Quota uint64 `json:"quota"`
Nodes uint64 `json:"nodes"`
}
type S3QoSResponse struct {
ApiLimitConf map[string]*UserLimitConf `json:"user_limit_conf"` // api --> userLimitConf
Nodes uint64 `json:"nodes"`
}
func IsS3PutApi(api string) bool {
switch strings.ToLower(api) {
case "putobject", "copyobject", "uploadpart", "uploadpartcopy", "postobject":
return true
default:
return false
}
}

View File

@ -944,3 +944,11 @@ func (api *AdminAPI) DelBucketLifecycle(volume string) (err error) {
}
return
}
func (api *AdminAPI) GetS3QoSInfo() (data []byte, err error) {
var request = newAPIRequest(http.MethodGet, proto.S3QoSGet)
if data, err = api.mc.serveRequest(request); err != nil {
return
}
return
}

View File

@ -0,0 +1,90 @@
// Copyright 2023 The CubeFS Authors.
//
// 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 concurrent
import (
"errors"
"sync"
"sync/atomic"
)
var (
ErrLimit = errors.New("limit exceeded")
)
type KeyConcurrentLimit struct {
mutex sync.RWMutex
current map[string]*int64 // uid --> concurrency
}
func NewLimit() *KeyConcurrentLimit {
return &KeyConcurrentLimit{current: make(map[string]*int64)}
}
func (s *KeyConcurrentLimit) Acquire(key string, limit int64) error {
s.mutex.RLock()
count := s.current[key]
s.mutex.RUnlock()
if count == nil {
s.mutex.Lock()
count = s.current[key]
if count == nil {
count = new(int64)
s.current[key] = count
}
s.mutex.Unlock()
}
newVal := atomic.AddInt64(count, 1)
if newVal > limit && limit != 0 {
atomic.AddInt64(count, -1)
return ErrLimit
}
return nil
}
func (s *KeyConcurrentLimit) Release(key string) {
s.mutex.RLock()
count, ok := s.current[key]
s.mutex.RUnlock()
if !ok {
return
}
val := atomic.AddInt64(count, -1)
if val < 0 {
atomic.StoreInt64(count, 0)
}
}
func (s *KeyConcurrentLimit) Running() int {
s.mutex.RLock()
defer s.mutex.RUnlock()
var all int64
for _, count := range s.current {
all += atomic.LoadInt64(count)
}
return int(all)
}
func (s *KeyConcurrentLimit) Get(key string) int64 {
s.mutex.RLock()
count, ok := s.current[key]
s.mutex.RUnlock()
if !ok {
return 0
}
return atomic.LoadInt64(count)
}

View File

@ -0,0 +1,90 @@
// Copyright 2023 The CubeFS Authors.
//
// 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 concurrent
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestSyncKeyLimit(t *testing.T) {
l := NewLimit()
key1, key2, key3 := string("1"), string("2"), string("3")
require.Equal(t, 0, l.Running())
require.NoError(t, l.Acquire(key1, 2))
require.NoError(t, l.Acquire(key2, 2))
require.NoError(t, l.Acquire(key1, 2))
require.NoError(t, l.Acquire(key2, 2))
require.NoError(t, l.Acquire(key3, 2))
require.Equal(t, ErrLimit, l.Acquire(key2, 2))
require.NoError(t, l.Acquire(key2, 3))
l.Release(key2)
require.Equal(t, 5, l.Running())
require.Equal(t, ErrLimit, l.Acquire(key1, 2))
require.Equal(t, ErrLimit, l.Acquire(key2, 2))
require.NoError(t, l.Acquire(key3, 2))
require.Equal(t, 6, l.Running())
l.Release(key1)
l.Release(key2)
require.Equal(t, 4, l.Running())
require.NoError(t, l.Acquire(key1, 2))
require.NoError(t, l.Acquire(key2, 2))
require.Equal(t, ErrLimit, l.Acquire(key3, 2))
require.Equal(t, 6, l.Running())
}
func TestKeyLimit_ReleaseNonExistKey(t *testing.T) {
ast := require.New(t)
key := "1"
l := NewLimit()
// should not panic
l.Release(key)
ast.Equal(int64(0), l.Get(key))
// should not decrease to minus integer
l.Acquire(key, 10)
ast.Equal(int64(1), l.Get(key))
for i := 0; i < 100; i++ {
l.Release(key)
ast.Equal(int64(0), l.Get(key))
}
}
func TestKeyLimit_Switch(t *testing.T) {
ast := require.New(t)
l := NewLimit()
key := "1"
ast.NoError(l.Acquire(key, 1))
ast.Error(l.Acquire(key, 1))
ast.Equal(int64(1), l.Get(key))
ast.NoError(l.Acquire(key, 0))
ast.Equal(int64(2), l.Get(key))
ast.Error(l.Acquire(key, 1))
ast.Equal(int64(2), l.Get(key))
ast.NoError(l.Acquire(key, 10))
ast.Equal(int64(3), l.Get(key))
}

122
util/flowctrl/controller.go Executable file
View File

@ -0,0 +1,122 @@
// Copyright 2023 The CubeFS Authors.
//
// 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 ratelimit RateLimit Algorithm Based on Token Bucket
package flowctrl
import (
"io"
"sync"
"time"
)
const (
tokenWindow = 50 * time.Millisecond
minWindowCap = 64
)
type Controller struct {
capacity int
threshold int
cond *sync.Cond
done chan struct{}
ratePerSecond int
closeOnce *sync.Once
}
func NewController(ratePerSecond int) *Controller {
capacity := ratePerSecond * int(tokenWindow) / int(time.Second)
if capacity < minWindowCap {
capacity = minWindowCap
}
self := &Controller{
ratePerSecond: ratePerSecond,
threshold: capacity,
capacity: capacity,
cond: sync.NewCond(new(sync.Mutex)),
done: make(chan struct{}, 1),
closeOnce: &sync.Once{},
}
// produce token asynchronously
go self.run(capacity)
return self
}
func (self *Controller) acquire(size int) int {
self.cond.L.Lock()
for self.capacity == 0 {
self.cond.Wait()
}
if size > self.capacity {
size = self.capacity
}
self.capacity -= size
self.cond.L.Unlock()
return size
}
func (self *Controller) fill(size int) {
if size <= 0 {
return
}
self.cond.L.Lock()
self.capacity += size
if self.capacity > self.threshold {
self.capacity = self.threshold
}
self.cond.L.Unlock()
self.cond.Broadcast()
}
func (self *Controller) run(capacity int) {
t := time.NewTicker(tokenWindow)
for {
select {
case <-t.C:
self.cond.L.Lock()
self.capacity = capacity
self.cond.L.Unlock()
self.cond.Broadcast()
case <-self.done:
t.Stop()
return
}
}
}
// Close release token-producing goroutine
func (self *Controller) Close() error {
self.closeOnce.Do(func() {
self.done <- struct{}{}
})
return nil
}
func (self *Controller) Reader(underlying io.Reader) io.Reader {
return &rateReader{
underlying: underlying,
c: self,
}
}
func (self *Controller) Writer(underlying io.Writer) io.Writer {
return &rateWriter{
underlying: underlying,
c: self,
}
}
func (self *Controller) GetRateLimit() int {
return self.ratePerSecond
}

115
util/flowctrl/controller_test.go Executable file
View File

@ -0,0 +1,115 @@
// Copyright 2023 The CubeFS Authors.
//
// 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 flowctrl
import (
"bytes"
"crypto/rand"
"io"
"log"
"math"
"runtime"
"sync"
"testing"
"time"
"github.com/stretchr/testify/assert"
)
var (
size = 4 * 1024 * 1024
b = make([]byte, size)
)
func init() {
runtime.GOMAXPROCS(4)
log.SetFlags(log.LstdFlags | log.Lshortfile)
rand.Read(b)
}
func TestControllerReader(t *testing.T) {
{
size := size
c := NewController(2 * 1024 * 1024)
defer c.Close()
assert.Equal(t, 2*1024*1024, c.GetRateLimit())
r1 := c.Reader(bytes.NewReader(b))
r2 := c.Reader(bytes.NewReader(b))
b1 := make([]byte, size)
b2 := make([]byte, size)
wg := sync.WaitGroup{}
wg.Add(2)
go func() {
n, err := io.ReadFull(r1, b1)
assert.NoError(t, err)
assert.Equal(t, size, n)
wg.Done()
}()
go func() {
time.Sleep(1 * time.Second)
n, err := io.ReadFull(r2, b2)
assert.NoError(t, err)
assert.Equal(t, size, n)
wg.Done()
}()
now := time.Now()
wg.Wait()
elapsed := time.Since(now).Seconds()
log.Println(elapsed)
assert.True(t, math.Abs(4-elapsed) < 0.5)
assert.Equal(t, b, b1)
assert.Equal(t, b, b2)
}
}
func TestControllerWriter(t *testing.T) {
{
size := int64(size)
c := NewController(2 * 1024 * 1024)
defer c.Close()
b1 := new(bytes.Buffer)
b2 := new(bytes.Buffer)
w1 := c.Writer(b1)
w2 := c.Writer(b2)
wg := sync.WaitGroup{}
wg.Add(2)
go func() {
n, err := io.Copy(w1, bytes.NewReader(b))
assert.NoError(t, err)
assert.Equal(t, size, n)
wg.Done()
}()
go func() {
time.Sleep(1 * time.Second)
n, err := io.Copy(w2, bytes.NewReader(b))
assert.NoError(t, err)
assert.Equal(t, size, n)
wg.Done()
}()
now := time.Now()
wg.Wait()
elapsed := time.Since(now).Seconds()
log.Println(elapsed)
assert.True(t, math.Abs(4-elapsed) < 0.5)
assert.Equal(t, b, b1.Bytes())
assert.Equal(t, b, b2.Bytes())
}
}

71
util/flowctrl/keycontroller.go Executable file
View File

@ -0,0 +1,71 @@
// Copyright 2023 The CubeFS Authors.
//
// 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 flowctrl
import (
"sync"
)
type ctrlWrapper struct {
c *Controller
refCount int
}
func newCtrlWrapper(rate int) *ctrlWrapper {
return &ctrlWrapper{
c: NewController(rate),
refCount: 0,
}
}
type KeyFlowCtrl struct {
mutex sync.RWMutex
current map[string]*ctrlWrapper // uid -> controller
}
func NewKeyFlowCtrl() *KeyFlowCtrl {
return &KeyFlowCtrl{current: make(map[string]*ctrlWrapper)}
}
func (k *KeyFlowCtrl) Acquire(key string, rate int) *Controller {
k.mutex.Lock()
ctrl, ok := k.current[key]
if !ok {
ctrl = newCtrlWrapper(rate)
k.current[key] = ctrl
}
ctrl.refCount++
k.mutex.Unlock()
return ctrl.c
}
func (k *KeyFlowCtrl) Release(key string) {
k.mutex.Lock()
ctrl, ok := k.current[key]
if !ok {
panic("key not in map. Possible reason: Release without Acquire.")
}
ctrl.refCount--
if ctrl.refCount < 0 {
panic("internal error: refs < 0")
}
if ctrl.refCount == 0 {
ctrl.c.Close() // avoid goroutine leak
delete(k.current, key)
}
k.mutex.Unlock()
}

View File

@ -0,0 +1,101 @@
// Copyright 2023 The CubeFS Authors.
//
// 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 flowctrl
import (
"fmt"
"io"
"io/ioutil"
"math/rand"
"strings"
"testing"
"github.com/stretchr/testify/assert"
)
func TestKeyController(t *testing.T) {
k := NewKeyFlowCtrl()
key1, key2, key3 := "10001", "10002", "10003"
k.Acquire(key1, 1024)
assert.Equal(t, 1, k.current[key1].refCount)
k.Acquire(key1, 1024)
assert.Equal(t, 2, k.current[key1].refCount)
k.Acquire(key2, 1024)
assert.Equal(t, 2, len(k.current))
k.Acquire(key3, 1024)
assert.Equal(t, 3, len(k.current))
k.Release(key1)
assert.Equal(t, 1, k.current[key1].refCount)
assert.Equal(t, 3, len(k.current))
k.Release(key1)
_, ok := k.current[key1]
assert.Equal(t, false, ok)
assert.Equal(t, 2, len(k.current))
k.Release(key2)
assert.Equal(t, 1, len(k.current))
k.Release(key3)
assert.Equal(t, 0, len(k.current))
assert.Panics(t, func() {
k.Release(key3)
})
}
func BenchmarkRWKeyRateCtrl(b *testing.B) {
k := NewKeyFlowCtrl()
b.SetParallelism(100)
b.RunParallel(func(pb *testing.PB) {
for pb.Next() {
uid := rand.Int31n(10000)
uidStr := fmt.Sprint(uid)
k.Acquire(uidStr, 1024)
defer k.Release(uidStr)
}
})
t := testing.T{}
assert.Equal(&t, 0, len(k.current))
}
func BenchmarkRWKeyRateCtrlReader(b *testing.B) {
str := initSource(8 * 1024)
k := NewKeyFlowCtrl()
b.SetParallelism(100)
b.RunParallel(func(pb *testing.PB) {
for pb.Next() {
uid := rand.Int31n(10000)
uidStr := fmt.Sprint(uid)
c := k.Acquire(uidStr, 1024)
defer k.Release(uidStr)
reader := NewRateReaderWithCtrl(strings.NewReader(str), c)
io.Copy(ioutil.Discard, reader)
}
})
t := testing.T{}
assert.Equal(&t, 0, len(k.current))
}
func initSource(size int) string {
str := "0123456789abcdefghijklmnopqrstuvwxyz"
res := make([]byte, 0, size)
for i := 0; i < size; i++ {
res = append(res, str[rand.Intn(len(str))])
}
return string(res)
}

49
util/flowctrl/reader.go Executable file
View File

@ -0,0 +1,49 @@
// Copyright 2023 The CubeFS Authors.
//
// 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 flowctrl
import "io"
type rateReader struct {
underlying io.Reader
c *Controller
}
func (self *rateReader) Read(p []byte) (n int, err error) {
size := len(p)
size = self.c.acquire(size)
n, err = self.underlying.Read(p[:size])
self.c.fill(size - n)
return
}
func NewRateReader(r io.Reader, ratePerSecond int) io.ReadCloser {
c := NewController(ratePerSecond)
r = c.Reader(r)
return struct {
io.Reader
io.Closer
}{
Reader: r,
Closer: c,
}
}
func NewRateReaderWithCtrl(r io.Reader, c *Controller) io.Reader {
r = c.Reader(r)
return r
}

96
util/flowctrl/reader_test.go Executable file
View File

@ -0,0 +1,96 @@
// Copyright 2023 The CubeFS Authors.
//
// 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 flowctrl
import (
"bytes"
"io"
"log"
"math"
"testing"
"testing/iotest"
"time"
"github.com/stretchr/testify/assert"
)
func TestRateReader(t *testing.T) {
{
size := size
rc := NewRateReader(bytes.NewReader(b), 1*1024*1024)
defer rc.Close()
b2 := make([]byte, size)
now := time.Now()
n, err := io.ReadFull(rc, b2)
elapsed := time.Since(now).Seconds()
log.Println(elapsed)
assert.True(t, math.Abs(4-elapsed) < 0.5)
assert.NoError(t, err)
assert.Equal(t, size, n)
assert.Equal(t, b, b2)
}
{
size := int64(size)
rc := NewRateReader(bytes.NewReader(b), 1*1024*1024)
defer rc.Close()
b2 := new(bytes.Buffer)
now := time.Now()
n, err := io.Copy(b2, rc)
elapsed := time.Since(now).Seconds()
log.Println(elapsed)
assert.True(t, math.Abs(4-elapsed) < 0.5)
assert.NoError(t, err)
assert.Equal(t, size, n)
assert.Equal(t, b, b2.Bytes())
}
{
size := size
rc := NewRateReader(iotest.TimeoutReader(bytes.NewReader(b)), 1*1024*1024)
defer rc.Close()
b2 := make([]byte, size)
n, err := io.ReadFull(rc, b2)
assert.Equal(t, iotest.ErrTimeout, err)
assert.Equal(t, 1*1024*1024*int(tokenWindow)/int(time.Second), n)
assert.Equal(t, b[:n], b2[:n])
}
{
size := 10
rc := NewRateReader(iotest.OneByteReader(bytes.NewReader(b)), 1*1024*1024)
defer rc.Close()
b2 := make([]byte, size)
n, err := io.ReadFull(rc, b2)
assert.NoError(t, err)
assert.Equal(t, size, n)
assert.Equal(t, b[:n], b2[:n])
}
{
size := 10
rc := NewRateReader(iotest.OneByteReader(bytes.NewReader(b[:size])), 4)
defer rc.Close()
b2 := make([]byte, size+1)
n, err := io.ReadFull(rc, b2)
assert.Equal(t, io.ErrUnexpectedEOF, err)
assert.Equal(t, size, n)
assert.Equal(t, b[:n], b2[:n])
}
}

57
util/flowctrl/writer.go Executable file
View File

@ -0,0 +1,57 @@
// Copyright 2023 The CubeFS Authors.
//
// 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 flowctrl
import "io"
type rateWriter struct {
underlying io.Writer
c *Controller
}
func (self *rateWriter) Write(p []byte) (written int, err error) {
write:
size := len(p)
size = self.c.acquire(size)
n, err := self.underlying.Write(p[:size])
self.c.fill(size - n)
written += n
if err != nil {
return
}
if size != len(p) {
p = p[size:]
goto write
}
return
}
func NewRateWriter(w io.Writer, ratePerSecond int) io.WriteCloser {
c := NewController(ratePerSecond)
w = c.Writer(w)
return struct {
io.Writer
io.Closer
}{
Writer: w,
Closer: c,
}
}
func NewRateWriterWithCtrl(w io.Writer, c *Controller) io.Writer {
w = c.Writer(w)
return w
}

81
util/flowctrl/writer_test.go Executable file
View File

@ -0,0 +1,81 @@
// Copyright 2023 The CubeFS Authors.
//
// 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 flowctrl
import (
"bytes"
"io"
"log"
"math"
"testing"
"testing/iotest"
"time"
"github.com/stretchr/testify/assert"
)
func TestRateWriter(t *testing.T) {
{
size := size
w := bytes.NewBuffer(nil)
wc := NewRateWriter(w, 1*1024*1024)
defer wc.Close()
now := time.Now()
n, err := wc.Write(b)
elapsed := time.Since(now).Seconds()
log.Println(elapsed)
assert.True(t, math.Abs(4-elapsed) < 0.5)
assert.NoError(t, err)
assert.Equal(t, size, n)
assert.Equal(t, b, w.Bytes())
}
{
size := int64(size)
w := bytes.NewBuffer(nil)
wc := NewRateWriter(w, 1*1024*1024)
defer wc.Close()
now := time.Now()
n, err := io.Copy(wc, bytes.NewBuffer(b))
elapsed := time.Since(now).Seconds()
log.Println(elapsed)
assert.True(t, math.Abs(4-elapsed) < 0.5)
assert.NoError(t, err)
assert.Equal(t, size, n)
assert.Equal(t, b, w.Bytes())
}
{
w := bytes.NewBuffer(nil)
wc := NewRateWriter(w, 5*1024)
defer wc.Close()
n, err := io.Copy(wc, iotest.TimeoutReader(bytes.NewReader(b)))
assert.Equal(t, iotest.ErrTimeout, err)
assert.Equal(t, b[:n], w.Bytes()[:n])
}
{
w := bytes.NewBuffer(nil)
wc := NewRateWriter(w, 5*1024)
defer wc.Close()
n, err := wc.Write(b[:1])
assert.NoError(t, err)
assert.Equal(t, 1, n)
assert.Equal(t, b[:1], w.Bytes())
}
}

View File

@ -0,0 +1,73 @@
// Copyright 2023 The CubeFS Authors.
//
// 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 ratelimit
import (
"sync"
"time"
"gopkg.in/bsm/ratelimit.v1"
)
type rateLimitWrapper struct {
r *ratelimit.RateLimiter
refCount int
}
func newRateLimitWrapper(rate uint64) *rateLimitWrapper {
return &rateLimitWrapper{
r: ratelimit.New(10*rate, 10*time.Second),
refCount: 0,
}
}
type KeyRateLimit struct {
mutex sync.RWMutex
current map[string]*rateLimitWrapper // uid -> ratelimit
}
func NewKeyRateLimit() *KeyRateLimit {
return &KeyRateLimit{current: make(map[string]*rateLimitWrapper)}
}
func (k *KeyRateLimit) Acquire(key string, rate uint64) *ratelimit.RateLimiter {
k.mutex.Lock()
limit, ok := k.current[key]
if !ok {
limit = newRateLimitWrapper(rate)
k.current[key] = limit
}
limit.refCount++
k.mutex.Unlock()
return limit.r
}
func (k *KeyRateLimit) Release(key string) {
k.mutex.Lock()
limit, ok := k.current[key]
if !ok {
panic("key not in map. Possible reason: Release without Acquire.")
}
limit.refCount--
if limit.refCount < 0 {
panic("internal error: refs < 0")
}
if limit.refCount == 0 {
delete(k.current, key)
}
k.mutex.Unlock()
}

View File

@ -0,0 +1,53 @@
// Copyright 2023 The CubeFS Authors.
//
// 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 ratelimit
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestKeyController(t *testing.T) {
k := NewKeyRateLimit()
key1, key2, key3 := "10001", "10002", "10003"
k.Acquire(key1, 1024)
assert.Equal(t, 1, k.current[key1].refCount)
k.Acquire(key1, 1024)
assert.Equal(t, 2, k.current[key1].refCount)
k.Acquire(key2, 1024)
assert.Equal(t, 2, len(k.current))
k.Acquire(key3, 1024)
assert.Equal(t, 3, len(k.current))
k.Release(key1)
assert.Equal(t, 1, k.current[key1].refCount)
assert.Equal(t, 3, len(k.current))
k.Release(key1)
_, ok := k.current[key1]
assert.Equal(t, false, ok)
assert.Equal(t, 2, len(k.current))
k.Release(key2)
assert.Equal(t, 1, len(k.current))
k.Release(key3)
assert.Equal(t, 0, len(k.current))
assert.Panics(t, func() {
k.Release(key3)
})
}

View File

@ -0,0 +1,168 @@
// Copyright 2023 The CubeFS Authors.
//
// 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 reloadconf
import (
"bytes"
"crypto/md5"
"encoding/base64"
"fmt"
"io/ioutil"
"os"
"sync"
"time"
"github.com/cubefs/cubefs/blobstore/util/errors"
"github.com/cubefs/cubefs/util/log"
)
// ReloadConf loads remote configuration change dynamically
type ReloadConf struct {
ConfName string
ReloadMs int
RequestRemote func() ([]byte, error)
md5sum []byte
mutex sync.Mutex
}
func (self *ReloadConf) reload(reload func(data []byte) error) error {
self.mutex.Lock()
defer self.mutex.Unlock()
// remote load
if self.RequestRemote != nil {
errRemote := self.remoteReload(reload)
if errRemote != nil {
log.LogWarn("remoteReload failed", self.ConfName, errors.Detail(errRemote))
} else {
return nil
}
}
// local load
errLocal := self.localReload(reload)
if errLocal != nil {
err := errors.Info(errLocal, "localReload").Detail(errLocal)
return err
}
return nil
}
func (self *ReloadConf) remoteReload(reload func(data []byte) error) (err error) {
data, md5sum, err := fetchRemote(self.RequestRemote)
if err != nil {
err = errors.Info(err, "fetchRemote").Detail(err)
return
}
// compare whether md5sum is changed
if bytes.Equal(md5sum, self.md5sum) {
log.LogDebug("remoteReload:", self.ConfName, "do nothing cause of md5sum is equal")
return nil
}
confName := fmt.Sprintf("%v_%v", self.ConfName, base64.URLEncoding.EncodeToString(md5sum))
err = ioutil.WriteFile(confName, data, 0666)
if err != nil {
err = errors.Info(err, "ioutil.WriteFile")
return
}
log.LogInfof("remoteReload %v remote file is changed, oldmd5: %v, newmd5: %v", self.ConfName, self.md5sum, md5sum)
err = reload(data)
if err != nil {
os.Remove(confName)
err = errors.Info(err, "reload", confName).Detail(err)
return
}
err = os.Rename(confName, self.ConfName)
if err != nil {
os.Remove(confName)
err = errors.Info(err, "os.Rename")
return
}
self.md5sum = md5sum
return
}
func (self *ReloadConf) localReload(reload func(data []byte) error) (err error) {
data, err := ioutil.ReadFile(self.ConfName)
if err != nil {
err = errors.Info(err, "ioutil.ReadFile").Detail(err)
return
}
md5sum := calcMD5Sum(data)
if bytes.Equal(md5sum, self.md5sum) {
log.LogDebug("localReload:", self.ConfName, "do nothing cause of md5sum is equal")
return nil
}
log.LogInfof("localReload: %v local file is changed, oldmd5: %v, newmd5: %v", self.ConfName, self.md5sum, md5sum)
err = reload(data)
if err != nil {
err = errors.Info(err, "reload").Detail(err)
return
}
self.md5sum = md5sum
return
}
func fetchRemote(requestRemote func() ([]byte, error)) (data, md5sum []byte, err error) {
data, err = requestRemote()
if err != nil {
err = errors.Info(err, "ioutil.ReadAll")
return
}
md5sum = calcMD5Sum(data)
return
}
func StartReload(cfg *ReloadConf, reload func(data []byte) error) (err error) {
err = cfg.reload(reload)
if err != nil {
log.LogError("cfg.reload:", cfg.ConfName, errors.Detail(err))
return
}
if cfg.ReloadMs == 0 {
return
}
go func() {
dur := time.Duration(cfg.ReloadMs) * time.Millisecond
for range time.Tick(dur) {
err := cfg.reload(reload)
if err != nil {
log.LogError("cfg.reload:", cfg.ConfName, errors.Detail(err))
}
}
}()
return
}
func calcMD5Sum(b []byte) []byte {
h := md5.New()
h.Write(b)
return h.Sum(nil)
}

13
vendor/gopkg.in/bsm/ratelimit.v1/Makefile generated vendored Executable file
View File

@ -0,0 +1,13 @@
default: test
testdeps:
@go get github.com/onsi/ginkgo
@go get github.com/onsi/gomega
test: testdeps
@go test ./...
testrace: testdeps
@go test ./... -race
testall: test testrace

54
vendor/gopkg.in/bsm/ratelimit.v1/README.md generated vendored Executable file
View File

@ -0,0 +1,54 @@
# RateLimit [![Build Status](https://travis-ci.org/bsm/ratelimit.png?branch=master)](https://travis-ci.org/bsm/ratelimit)
Simple, thread-safe Go rate-limiter.
Inspired by Antti Huima's algorithm on http://stackoverflow.com/a/668327
### Example
```go
package main
import (
"github.com/bsm/redeo"
"log"
)
func main() {
// Create a new rate-limiter, allowing up-to 10 calls
// per second
rl := ratelimit.New(10, time.Second)
for i:=0; i<20; i++ {
if rl.Limit() {
fmt.Println("DOH! Over limit!")
} else {
fmt.Println("OK")
}
}
}
```
### Licence
```
Copyright (c) 2015 Black Square Media
Permission is hereby granted, free of charge, to any person obtaining
a copy of this software and associated documentation files (the
"Software"), to deal in the Software without restriction, including
without limitation the rights to use, copy, modify, merge, publish,
distribute, sublicense, and/or sell copies of the Software, and to
permit persons to whom the Software is furnished to do so, subject to
the following conditions:
The above copyright notice and this permission notice shall be
included in all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
```

89
vendor/gopkg.in/bsm/ratelimit.v1/ratelimit.go generated vendored Executable file
View File

@ -0,0 +1,89 @@
/*
Simple, thread-safe Go rate-limiter.
Inspired by Antti Huima's algorithm on http://stackoverflow.com/a/668327
Example:
// Create a new rate-limiter, allowing up-to 10 calls
// per second
rl := ratelimit.New(10, time.Second)
for i:=0; i<20; i++ {
if rl.Limit() {
fmt.Println("DOH! Over limit!")
} else {
fmt.Println("OK")
}
}
*/
package ratelimit
import (
"sync/atomic"
"time"
)
// RateLimit instances are thread-safe.
type RateLimiter struct {
rate, allowance, max, unit, lastCheck uint64
}
// New creates a new rate limiter instance
func New(rate uint64, per time.Duration) *RateLimiter {
nano := uint64(per)
if nano < 1 {
nano = uint64(time.Second)
}
if rate < 1 {
rate = 1
}
return &RateLimiter{
rate: uint64(rate), // store the rate
allowance: uint64(rate) * nano, // set our allowance to max in the beginning
max: uint64(rate) * nano, // remember our maximum allowance
unit: nano, // remember our unit size
lastCheck: unixNano(),
}
}
// Limit returns true if rate was exceeded
func (rl *RateLimiter) Limit() bool {
// Calculate the number of ns that have passed since our last call
now := unixNano()
passed := now - atomic.SwapUint64(&rl.lastCheck, now)
// Add them to our allowance
current := atomic.AddUint64(&rl.allowance, passed*rl.rate)
// Ensure our allowance is not over maximum
if current > rl.max {
atomic.AddUint64(&rl.allowance, rl.max-current)
current = rl.max
}
// If our allowance is less than one unit, rate-limit!
if current < rl.unit {
return true
}
// Not limited, subtract a unit
atomic.AddUint64(&rl.allowance, -rl.unit)
return false
}
// Undo reverts the last Limit() call, returning consumed allowance
func (rl *RateLimiter) Undo() {
current := atomic.AddUint64(&rl.allowance, rl.unit)
// Ensure our allowance is not over maximum
if current > rl.max {
atomic.AddUint64(&rl.allowance, rl.max-current)
}
}
// now as unix nanoseconds
func unixNano() uint64 {
return uint64(time.Now().UnixNano())
}

106
vendor/gopkg.in/bsm/ratelimit.v1/ratelimit_test.go generated vendored Executable file
View File

@ -0,0 +1,106 @@
package ratelimit
import (
"sync"
"testing"
"time"
. "github.com/onsi/ginkgo"
. "github.com/onsi/gomega"
)
var _ = Describe("RateLimiter", func() {
It("should accurately rate-limit at small rates", func() {
var count int
rl := New(10, time.Minute)
for !rl.Limit() {
count++
}
Expect(count).To(Equal(10))
})
It("should accurately rate-limit at large rates", func() {
var count int
rl := New(100000, time.Hour)
for !rl.Limit() {
count++
}
Expect(count).To(BeNumerically("~", 100000, 10))
})
It("should accurately rate-limit at large intervals", func() {
var count int
rl := New(100, 360*24*time.Hour)
for !rl.Limit() {
count++
}
Expect(count).To(Equal(100))
})
It("should correctly increase allowance", func() {
n := 25
rl := New(n, 50*time.Millisecond)
for i := 0; i < n; i++ {
Expect(rl.Limit()).To(BeFalse(), "on cycle %d", i)
}
Expect(rl.Limit()).To(BeTrue())
Eventually(rl.Limit, "60ms", "10ms").Should(BeFalse())
})
It("should correctly spread allowance", func() {
var count int
rl := New(5, 10*time.Millisecond)
start := time.Now()
for time.Now().Sub(start) < 100*time.Millisecond {
if !rl.Limit() {
count++
}
}
Expect(count).To(BeNumerically("~", 54, 1))
})
It("should undo", func() {
rl := New(5, time.Minute)
Expect(rl.Limit()).To(BeFalse())
Expect(rl.Limit()).To(BeFalse())
Expect(rl.Limit()).To(BeFalse())
Expect(rl.Limit()).To(BeFalse())
Expect(rl.Limit()).To(BeFalse())
Expect(rl.Limit()).To(BeTrue())
rl.Undo()
Expect(rl.Limit()).To(BeFalse())
Expect(rl.Limit()).To(BeTrue())
})
It("should be thread-safe", func() {
c := 100
n := 100
wg := sync.WaitGroup{}
rl := New(c*n, time.Hour)
for i := 0; i < c; i++ {
wg.Add(1)
go func(thread int) {
defer GinkgoRecover()
defer wg.Done()
for j := 0; j < n; j++ {
Expect(rl.Limit()).To(BeFalse(), "thread %d, cycle %d", thread, j)
}
}(i)
}
wg.Wait()
Expect(rl.Limit()).To(BeTrue())
})
})
// --------------------------------------------------------------------
func TestGinkgoSuite(t *testing.T) {
RegisterFailHandler(Fail)
RunSpecs(t, "github.com/bsm/ratelimit")
}