Skip to content

Commit

Permalink
Fix startup procedure: Fixes potential ghost death and console handler
Browse files Browse the repository at this point in the history
  • Loading branch information
tterrag1098 committed Apr 29, 2021
1 parent f074a08 commit 02c4e91
Show file tree
Hide file tree
Showing 2 changed files with 110 additions and 100 deletions.
19 changes: 13 additions & 6 deletions src/main/java/com/tterrag/k9/K9.java
Original file line number Diff line number Diff line change
Expand Up @@ -130,8 +130,8 @@ public Mono<Void> start() {
GatewayBootstrap<GatewayOptions> gateway = client.gateway()
.setEventDispatcher(ReplayingEventDispatcher.builder()
.replayEventFilter(e -> e instanceof ReadyEvent)
.eventScheduler(Schedulers.boundedElastic())
.build())
.eventScheduler(Schedulers.boundedElastic())
.build())
.setEnabledIntents(IntentSet.of(
Intent.GUILDS, Intent.GUILD_MEMBERS, Intent.GUILD_PRESENCES,
Intent.GUILD_MESSAGES, Intent.GUILD_MESSAGE_REACTIONS,
Expand Down Expand Up @@ -193,8 +193,15 @@ public Mono<Void> start() {

return Mono.fromRunnable(commands::slurpCommands)
.then(gateway.login())
.flatMap(c -> Mono.when(onInitialReady.apply(c.getEventDispatcher()), services.start(c)).thenReturn(c))
.flatMap(this::teardown);
.flatMap(c ->
Mono.when(onInitialReady.apply(c.getEventDispatcher()), services.start(c))
.doOnError(t -> log.error("Unexpected error received in main bot subscriber:", t))
.doOnTerminate(() -> log.error("Unexpected completion of main bot subscriber!"))
.zipWith(teardown(c))
.thenReturn(c)
.onErrorResume($ -> teardown(c).thenReturn(c)))
.flatMap(c -> Mono.fromRunnable(commands::onShutdown)
.then(c.logout()));
}

private boolean isUser(MessageCreateEvent evt) {
Expand All @@ -210,7 +217,7 @@ private Mono<Void> teardown(GatewayDiscordClient gatewayClient) {
while (scan.hasNextLine()) {
if (scan.nextLine().equals("stop")) {
scan.close();
System.exit(0);
return null; // Empty completion will bubble up to zip below
}
}
Threads.sleep(100);
Expand All @@ -226,7 +233,7 @@ private Mono<Void> teardown(GatewayDiscordClient gatewayClient) {

return Mono.zip(consoleHandler, gatewayClient.onDisconnect())
.then()
.doOnTerminate(() -> log.error("Unexpected completion of main bot subscriber!"));
.doOnError(t -> log.error("Disconnect listener error:", t));
}

public static String getVersion() {
Expand Down
191 changes: 97 additions & 94 deletions src/main/java/com/tterrag/k9/commands/api/CommandRegistrar.java
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
Expand Down Expand Up @@ -41,59 +42,60 @@

@Slf4j
public class CommandRegistrar {
@NonNull
static final File DATA_FOLDER = NullHelper.notnullJ(Paths.get("command_data").toFile(), "Path#toFile");
static {
DATA_FOLDER.mkdirs();
}

private final K9 k9;

private final Map<String, ICommand> commands = Maps.newTreeMap();
private final CommandControl ctrl = new CommandControl();

private final @NonNull GsonBuilder builder = new GsonBuilder();
private @NonNull Gson gson = new Gson();

private boolean finishedDefaultSlurp;
private boolean locked;

private Disposable autoSaveSubscriber;

public CommandRegistrar(K9 k9) {
this.k9 = k9;
}
static final File DATA_FOLDER = NullHelper.notnullJ(Paths.get("command_data").toFile(), "Path#toFile");
static {
DATA_FOLDER.mkdirs();
}

private final K9 k9;

private final Map<String, ICommand> commands = Maps.newTreeMap();
private final CommandControl ctrl = new CommandControl();

private final @NonNull GsonBuilder builder = new GsonBuilder();
private @NonNull Gson gson = new Gson();

private boolean finishedDefaultSlurp;
private boolean locked;
private AtomicBoolean shutdown = new AtomicBoolean(false);

private Disposable autoSaveSubscriber;

public CommandRegistrar(K9 k9) {
this.k9 = k9;
}

public Mono<ICommand> invokeCommand(MessageCreateEvent evt, String name, String argstr) {
Optional<ICommand> commandReq = findCommand(evt.getGuildId().orElse(null), name);
ICommand command = commandReq.filter(c -> !c.admin() || evt.getMessage().getAuthor().map(this::isAdmin).orElse(false)).orElse(null);
if (command == null) {
return Mono.empty();
}
CommandContext ctx = new CommandContext(k9, evt);
if (!command.requirements().matches(ctx).block()) {
return evt.getMessage().getChannel()
.flatMap(c -> c.createMessage("You do not have permission to use this command!"))
.delayElement(Duration.ofSeconds(5))
.flatMap(m -> m.delete())
.thenReturn(command);
}
// evt.getMessage().getChannel().flatMap(c -> c.type()).subscribe();
argstr = Strings.nullToEmpty(argstr);
Map<Flag, String> flags = new HashMap<>();
Map<Argument<?>, String> args = new HashMap<>();
Map<Character, Flag> keyToFlag = command.getFlags().stream().collect(Collectors.toMap(Flag::name, f -> f));
Map<String, Flag> longKeyToFlag = command.getFlags().stream().collect(Collectors.toMap(Flag::longFormName, f -> f));
public Mono<ICommand> invokeCommand(MessageCreateEvent evt, String name, String argstr) {
Optional<ICommand> commandReq = findCommand(evt.getGuildId().orElse(null), name);
ICommand command = commandReq.filter(c -> !c.admin() || evt.getMessage().getAuthor().map(this::isAdmin).orElse(false)).orElse(null);
if (command == null) {
return Mono.empty();
}
CommandContext ctx = new CommandContext(k9, evt);
if (!command.requirements().matches(ctx).block()) {
return evt.getMessage().getChannel()
.flatMap(c -> c.createMessage("You do not have permission to use this command!"))
.delayElement(Duration.ofSeconds(5))
.flatMap(m -> m.delete())
.thenReturn(command);
}
// evt.getMessage().getChannel().flatMap(c -> c.type()).subscribe();
argstr = Strings.nullToEmpty(argstr);
Map<Flag, String> flags = new HashMap<>();
Map<Argument<?>, String> args = new HashMap<>();
Map<Character, Flag> keyToFlag = command.getFlags().stream().collect(Collectors.toMap(Flag::name, f -> f));
Map<String, Flag> longKeyToFlag = command.getFlags().stream().collect(Collectors.toMap(Flag::longFormName, f -> f));

Matcher matcher = Patterns.FLAGS.matcher(argstr);
Matcher matcher = Patterns.FLAGS.matcher(argstr);
while (matcher.find()) {
String flagname = matcher.group(2);
List<Flag> foundFlags;
Expand Down Expand Up @@ -165,14 +167,14 @@ public Mono<ICommand> invokeCommand(MessageCreateEvent evt, String name, String
return ctx.reply("Unexpected error processing command: " + e).thenReturn(command); // TODO should this be different?
}
}
public boolean isAdmin(User user) {
return k9.isAdmin(user.getId());
}
public Optional<ICommand> findCommand(CommandContext ctx, String name) {
return findCommand(ctx.getGuildId().orElse(null), name);
}
public boolean isAdmin(User user) {
return k9.isAdmin(user.getId());
}
public Optional<ICommand> findCommand(CommandContext ctx, String name) {
return findCommand(ctx.getGuildId().orElse(null), name);
}

public Optional<ICommand> findCommand(@Nullable Snowflake guild, String name) {
if (guild != null && ctrl.getData(guild).getCommandBlacklist().contains(name)) {
Expand All @@ -188,7 +190,7 @@ public void slurpCommands() {
}
}

@SneakyThrows
@SneakyThrows
public void slurpCommands(@NonNull String packagename) {
if (locked) {
throw new IllegalStateException("Cannot slurp commands in locked registrar.");
Expand All @@ -199,28 +201,28 @@ public void slurpCommands(@NonNull String packagename) {
return; // ??
}
ClassPath classpath = ClassPath.from(loader);
for (ClassInfo foo : classpath.getTopLevelClassesRecursive(packagename)) {
if (!foo.getName().equals(getClass().getName())) {
Class<?> c = foo.load();
if (c.isAnnotationPresent(Command.class)) {
log.info("Found annotation command: {}", c.getName());
registerCommand((ICommand) c.newInstance());
}
}
}
}
public void registerCommand(ICommand command) {
if (locked) {
throw new IllegalStateException("Cannot register command to locked registrar.");
}
if (!command.isTransient()) {
commands.put(command.getName(), command);
command.gatherParsers(builder);
command.onRegister(k9);
}
command.getChildren().forEach(this::registerCommand);
}
for (ClassInfo foo : classpath.getTopLevelClassesRecursive(packagename)) {
if (!foo.getName().equals(getClass().getName())) {
Class<?> c = foo.load();
if (c.isAnnotationPresent(Command.class)) {
log.info("Found annotation command: {}", c.getName());
registerCommand((ICommand) c.newInstance());
}
}
}
}
public void registerCommand(ICommand command) {
if (locked) {
throw new IllegalStateException("Cannot register command to locked registrar.");
}
if (!command.isTransient()) {
commands.put(command.getName(), command);
command.gatherParsers(builder);
command.onRegister(k9);
}
command.getChildren().forEach(this::registerCommand);
}

public void unregisterCommand(ICommand command) {
commands.remove(command.getName());
Expand Down Expand Up @@ -249,19 +251,20 @@ private void saveAll() {
}
}

public void onShutdown() {
saveAll();
for (ICommand c : commands.values()) {
c.onShutdown();
}
if (autoSaveSubscriber != null) {
autoSaveSubscriber.dispose();
}
}

public Iterable<ICommand> getCommands(Optional<Snowflake> guild) {
return getCommands(guild.orElse(null));
}
public void onShutdown() {
if (shutdown.getAndSet(true)) return;
saveAll();
for (ICommand c : commands.values()) {
c.onShutdown();
}
if (autoSaveSubscriber != null) {
autoSaveSubscriber.dispose();
}
}

public Iterable<ICommand> getCommands(Optional<Snowflake> guild) {
return getCommands(guild.orElse(null));
}

public Iterable<ICommand> getCommands(@Nullable Snowflake guild) {
if (guild == null) {
Expand Down

0 comments on commit 02c4e91

Please sign in to comment.