forked from botlabs-gg/yagpdb
-
Notifications
You must be signed in to change notification settings - Fork 0
/
bot.go
242 lines (201 loc) · 6.74 KB
/
bot.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
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
package bot
import (
"github.com/jonas747/discordgo"
"github.com/jonas747/dshardmanager"
"github.com/jonas747/dshardorchestrator/node"
"github.com/jonas747/dstate"
"github.com/jonas747/yagpdb/bot/deletequeue"
"github.com/jonas747/yagpdb/bot/eventsystem"
"github.com/jonas747/yagpdb/common"
"github.com/mediocregopher/radix.v3"
log "github.com/sirupsen/logrus"
"os"
"strconv"
"sync"
"time"
)
var (
// When the bot was started
Started = time.Now()
Enabled bool // wether the bot is set to run at some point in this process
Running bool // wether the bot is currently running
State *dstate.State
ShardManager *dshardmanager.Manager
StateHandlerPtr *eventsystem.Handler
NodeConn *node.Conn
UsingOrchestrator bool
MessageDeleteQueue = deletequeue.NewQueue()
)
var (
// the variables below specify shard clustering information, for splitting shards across multiple processes and machines
// this is very work in process and does not work at all atm
//
// plugins that needs to perform shard specific tasks, not directly related to incoming gateway events should check here
// to make sure the action they're doing is relevant to the current cluster (example: scheduled events should only run events for guilds on this cluster)
// the total amount of shards this bot is set to use across all processes
totalShardCount int
// The shards running on this process, protected by the processShardsLock muted
processShards []int
processShardsLock sync.RWMutex
)
func setup() {
// Things may rely on state being available at this point for initialization
State = dstate.NewState()
State.MaxChannelMessages = 1000
State.MaxMessageAge = time.Hour
// State.Debug = true
State.ThrowAwayDMMessages = true
State.TrackPrivateChannels = false
State.CacheExpirey = time.Minute * 10
go State.RunGCWorker()
eventsystem.AddHandler(HandleReady, eventsystem.EventReady)
StateHandlerPtr = eventsystem.AddHandler(StateHandler, eventsystem.EventAll)
eventsystem.ConcurrentAfter = StateHandlerPtr
eventsystem.AddHandler(ConcurrentEventHandler(EventLogger.handleEvent), eventsystem.EventAll)
eventsystem.AddHandler(HandleGuildCreate, eventsystem.EventGuildCreate)
eventsystem.AddHandler(HandleGuildDelete, eventsystem.EventGuildDelete)
eventsystem.AddHandler(HandleGuildUpdate, eventsystem.EventGuildUpdate)
eventsystem.AddHandler(HandleGuildRoleCreate, eventsystem.EventGuildRoleCreate)
eventsystem.AddHandler(HandleGuildRoleUpdate, eventsystem.EventGuildRoleUpdate)
eventsystem.AddHandler(HandleGuildRoleRemove, eventsystem.EventGuildRoleDelete)
eventsystem.AddHandler(HandleChannelCreate, eventsystem.EventChannelCreate)
eventsystem.AddHandler(HandleChannelUpdate, eventsystem.EventChannelUpdate)
eventsystem.AddHandler(HandleChannelDelete, eventsystem.EventChannelDelete)
eventsystem.AddHandler(HandleGuildMemberUpdate, eventsystem.EventGuildMemberUpdate)
eventsystem.AddHandler(HandleGuildMemberAdd, eventsystem.EventGuildMemberAdd)
eventsystem.AddHandler(HandleGuildMembersChunk, eventsystem.EventGuildMembersChunk)
}
func Run() {
setup()
log.Println("Running bot")
connEvtChannel, _ := strconv.ParseInt(os.Getenv("YAGPDB_CONNEVT_CHANNEL"), 10, 64)
connStatusChannel, _ := strconv.ParseInt(os.Getenv("YAGPDB_CONNSTATUS_CHANNEL"), 10, 64)
// Set up shard manager
ShardManager = dshardmanager.New(common.Conf.BotToken)
ShardManager.LogChannel = connEvtChannel
ShardManager.StatusMessageChannel = connStatusChannel
ShardManager.Name = "YAGPDB"
ShardManager.GuildCountsFunc = GuildCountsFunc
ShardManager.SessionFunc = func(token string) (session *discordgo.Session, err error) {
session, err = discordgo.New(token)
if err != nil {
return
}
session.StateEnabled = false
session.LogLevel = discordgo.LogInformational
session.SyncEvents = true
// Certain discordgo internals expect this to be present
// but in case of shard migration it's not, so manually assign it here
session.State.Ready = discordgo.Ready{
User: &discordgo.SelfUser{
User: common.BotUser,
},
}
return
}
// Only handler
ShardManager.AddHandler(eventsystem.HandleEvent)
orcheStratorAddress := os.Getenv("YAGPDB_ORCHESTRATOR_ADDRESS")
if orcheStratorAddress != "" {
UsingOrchestrator = true
log.Infof("Set to use orchestrator at address: %s", orcheStratorAddress)
} else {
log.Info("Running standalone without any orchestrator")
SetupStandalone()
}
go MemberFetcher.Run()
go mergedMessageSender()
Running = true
if UsingOrchestrator {
// TODO
NodeConn = node.NewNodeConn(&NodeImpl{}, orcheStratorAddress, common.VERSION, nil)
NodeConn.Run()
} else {
go ShardManager.Start()
InitPlugins()
}
// if masterAddr != "" {
// stateLock.Lock()
// state = StateWaitingHelloMaster
// stateLock.Unlock()
// log.Println("Connecting to master at ", masterAddr, ", wont start until connected and told to start")
// var err error
// SlaveClient, err = slave.ConnectToMaster(&SlaveImpl{}, masterAddr)
// if err != nil {
// log.WithError(err).Error("Failed connecting to master")
// os.Exit(1)
// }
// } else {
// stateLock.Lock()
// state = StateRunningNoMaster
// stateLock.Unlock()
// InitPlugins()
// log.Println("Running normally without a master")
// go ShardManager.Start()
// go MonitorLoading()
// }
// for _, p := range common.Plugins {
// starter, ok := p.(BotStarterHandler)
// if ok {
// starter.StartBot()
// log.Debug("Ran StartBot for ", p.Name())
// }
// }
}
func SetupStandalone() {
shardCount, err := ShardManager.GetRecommendedCount()
if err != nil {
panic("Failed getting shard count: " + err.Error())
}
totalShardCount = shardCount
processShards = make([]int, totalShardCount)
for i := 0; i < totalShardCount; i++ {
processShards[i] = i
}
err = common.RedisPool.Do(radix.FlatCmd(nil, "SET", "yagpdb_total_shards", shardCount))
if err != nil {
log.WithError(err).Error("failed setting shard count")
}
}
func InitPlugins() {
// Initialize all plugins
for _, plugin := range common.Plugins {
if initBot, ok := plugin.(BotInitHandler); ok {
initBot.BotInit()
}
}
}
func BotStarted() {
go loopCheckAdmins()
}
var stopOnce sync.Once
func StopAllPlugins(wg *sync.WaitGroup) {
stopOnce.Do(func() {
for _, v := range common.Plugins {
stopper, ok := v.(BotStopperHandler)
if !ok {
continue
}
wg.Add(1)
log.Debug("Calling bot stopper for: ", v.Name())
go stopper.StopBot(wg)
}
close(stopRunCheckAdmins)
})
}
func Stop(wg *sync.WaitGroup) {
StopAllPlugins(wg)
ShardManager.StopAll()
wg.Done()
}
func GuildCountsFunc() []int {
numShards := ShardManager.GetNumShards()
result := make([]int, numShards)
State.RLock()
for _, v := range State.Guilds {
shard := (v.ID >> 22) % int64(numShards)
result[shard]++
}
State.RUnlock()
return result
}