-
Notifications
You must be signed in to change notification settings - Fork 263
/
limit_reader.go
91 lines (80 loc) · 2.44 KB
/
limit_reader.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
/*
* Copyright 2020 The Dragonfly 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 limitreader
import (
"fmt"
"hash"
"io"
"github.com/opencontainers/go-digest"
"d7y.io/dragonfly/v2/pkg/ratelimiter/ratelimiter"
"d7y.io/dragonfly/v2/pkg/util/digestutils"
)
// NewLimitReader creates a LimitReader.
// src: reader
// rate: bytes/second
func NewLimitReader(src io.Reader, rate int64) *LimitReader {
return NewLimitReaderWithLimiter(newRateLimiterWithDefaultWindow(rate), src)
}
// NewLimitReaderWithLimiter creates LimitReader with a rateLimiter.
// src: reader
// rate: bytes/second
func NewLimitReaderWithLimiter(rl *ratelimiter.RateLimiter, src io.Reader) *LimitReader {
return &LimitReader{
Src: src,
Limiter: rl,
}
}
func newRateLimiterWithDefaultWindow(rate int64) *ratelimiter.RateLimiter {
return ratelimiter.NewRateLimiter(ratelimiter.TransRate(rate), 2)
}
// LimitReader reads stream with RateLimiter.
type LimitReader struct {
Src io.Reader
Limiter *ratelimiter.RateLimiter
digest hash.Hash
digestType string
}
func (lr *LimitReader) Read(p []byte) (n int, err error) {
n, e := lr.Src.Read(p)
if e != nil && e != io.EOF {
return n, e
}
if n > 0 {
if lr.digest != nil {
lr.digest.Write(p[:n])
}
lr.Limiter.AcquireBlocking(int64(n))
}
return n, e
}
// NewLimitReaderWithLimiterAndDigest creates LimitReader with rateLimiter and digest.
// src: reader
// rate: bytes/second
func NewLimitReaderWithLimiterAndDigest(src io.Reader, rl *ratelimiter.RateLimiter, digest hash.Hash, digestType digest.Algorithm) *LimitReader {
return &LimitReader{
Src: src,
Limiter: rl,
digest: digest,
digestType: digestType.String(),
}
}
// Digest calculates the digest of all contents read, return value is like <algo>:<hex_value>
func (lr *LimitReader) Digest() string {
if lr.digest != nil {
return fmt.Sprintf("%s:%s", lr.digestType, digestutils.ToHashString(lr.digest))
}
return ""
}