/
client.go
96 lines (75 loc) · 1.51 KB
/
client.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
92
93
94
95
96
/*
* Revision History:
* Initial: 2018/07/04 Li Zebang
*/
package pool
import (
"context"
"errors"
"strconv"
"sync"
"time"
"github.com/google/go-github/github"
"golang.org/x/oauth2"
)
const (
defaultRemaining = -1
// HeaderXRateLimitRemaining -
HeaderXRateLimitRemaining = "X-RateLimit-Remaining"
// HeaderXRateLimitReset -
HeaderXRateLimitReset = "X-RateLimit-Reset"
)
// Token -
type Token struct {
Tag string
Token string
}
// Client -
type Client struct {
*github.Client
isUsed bool
remaining int
lock *sync.Mutex
reset time.Time
tag string
}
func client(token *Token) (*Client, error) {
if token == nil {
return nil, errors.New("token cannot be nil")
}
ctx := context.Background()
ts := oauth2.StaticTokenSource(
&oauth2.Token{AccessToken: token.Token},
)
tc := oauth2.NewClient(ctx, ts)
return &Client{
Client: github.NewClient(tc),
lock: &sync.Mutex{},
remaining: defaultRemaining,
tag: token.Tag,
}, nil
}
// HandleResponse -
func (c *Client) HandleResponse(resp *github.Response) error {
c.lock.Lock()
defer c.lock.Unlock()
remaining, err := strconv.Atoi(resp.Header.Get(HeaderXRateLimitRemaining))
if err != nil {
return err
}
c.remaining = remaining
if c.isUsed {
return nil
}
reset, err := strconv.ParseInt(resp.Header.Get(HeaderXRateLimitReset), 10, 64)
if err != nil {
return err
}
c.reset = time.Unix(reset, 0)
return nil
}
// Reset -
func (c *Client) Reset() {
c.isUsed = false
c.remaining = defaultRemaining
}