forked from botlabs-gg/yagpdb
-
Notifications
You must be signed in to change notification settings - Fork 0
/
batchmemberjob.go
117 lines (96 loc) · 2.47 KB
/
batchmemberjob.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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
package bot
import (
"errors"
"github.com/jonas747/discordgo"
"github.com/jonas747/yagpdb/bot/eventsystem"
"sync"
"time"
)
var (
BatchMemberJobManager = newBatchMemberJobManager()
)
type batchMemberJob struct {
CreatedAt time.Time
GuildID int64
F func(guildID int64, members []*discordgo.Member)
targetCount int
numHandled int
lastHandledEvt time.Time
}
type batchMemberJobManager struct {
jobs []*batchMemberJob
mu sync.Mutex
}
func newBatchMemberJobManager() *batchMemberJobManager {
m := &batchMemberJobManager{}
go m.monitor()
return m
}
func (m *batchMemberJobManager) monitor() {
ticker := time.NewTicker(time.Second)
for {
<-ticker.C
m.checkall()
}
}
func (m *batchMemberJobManager) checkall() {
m.mu.Lock()
defer m.mu.Unlock()
OUTER:
for {
for i, v := range m.jobs {
inactiveFor := time.Since(v.lastHandledEvt)
if (v.numHandled > 0 && inactiveFor > time.Second*10) ||
(v.numHandled == 0 && inactiveFor > time.Second*60) ||
v.numHandled == v.targetCount {
m.jobs = append(m.jobs[:i], m.jobs[i+1:]...)
logger.Infof("batchMemberManager: Removed expired job on %d, handled: [%d/%d]", v.GuildID, v.numHandled, v.targetCount)
continue OUTER
}
}
break
}
}
var (
ErrGuildNotOnProcess = errors.New("Guild not on process")
)
func (m *batchMemberJobManager) NewBatchMemberJob(guildID int64, f func(guildID int64, member []*discordgo.Member)) error {
if !IsGuildOnCurrentProcess(guildID) {
return ErrGuildNotOnProcess
}
gs := State.Guild(true, guildID)
if gs == nil {
return ErrGuildNotFound
}
gs.RLock()
targetCount := gs.Guild.MemberCount
gs.RUnlock()
job := &batchMemberJob{
CreatedAt: time.Now(),
GuildID: guildID,
F: f,
lastHandledEvt: time.Now(),
targetCount: targetCount,
}
m.mu.Lock()
m.jobs = append(m.jobs, job)
m.mu.Unlock()
session := ShardManager.SessionForGuild(guildID)
if session == nil {
return errors.New("No session?")
}
session.GatewayManager.RequestGuildMembers(guildID, "", 0)
return nil
}
func (m *batchMemberJobManager) handleGuildMemberChunk(evt *eventsystem.EventData) {
chunk := evt.GuildMembersChunk()
m.mu.Lock()
defer m.mu.Unlock()
for _, v := range m.jobs {
if v.GuildID == chunk.GuildID {
go v.F(chunk.GuildID, chunk.Members)
v.numHandled += len(chunk.Members)
v.lastHandledEvt = time.Now()
}
}
}