From e6be2883020b43aa6035df5b2f40e3fb7db9367a Mon Sep 17 00:00:00 2001 From: jake Date: Thu, 25 Jun 2026 16:02:47 -0400 Subject: [PATCH] decoupling start and shutdown - start is no longer a blocking call - giving it back to users - the Stop method now waits for the timeout before closing --- README.md | 25 ++++++++++++++++-------- bolt.go | 57 +++++++++++++++++++++++-------------------------------- option.go | 10 ++++++++++ 3 files changed, 51 insertions(+), 41 deletions(-) diff --git a/README.md b/README.md index f868dbb..ba11613 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,8 @@ package main import ( "fmt" + "os" + "os/signal" "strings" "time" @@ -21,20 +23,20 @@ import ( /* A basic example of a bot with two commands and a general message handler for non-command messages. The bot uses a -verbose log level which will log everything for debugging purposes, and registers Discord Intents for message and -admin related permissions. This allows the bot to parse messages, send them, delete them, etc. as well as timeout +verbose log level which will log everything for debugging purposes, and registers Discord Intents for message and +admin related permissions. This allows the bot to parse messages, send them, delete them, etc. as well as timeout and mute users. This example registers three commands: 1. .ping - a basic ping/pong command that can be run by anyone at any time -2. .wait - a dummy command that replies "okay" it can only be run by users with the "user" role and can only be ran +2. .wait - a dummy command that replies "okay" it can only be run by users with the "user" role and can only be ran once every 25 seconds -3. .timeout - a admin command that can only be run by users with an "admin" role, this command will timeout any mentioned +3. .timeout - a admin command that can only be run by users with an "admin" role, this command will timeout any mentioned users for 5 minutes -A message handler is also registered in this example, message handlers are used to handle messages that do not contain a -command. This enables auto-moderation from the bot without manual intervention. The example message handler does two arbitrary +A message handler is also registered in this example, message handlers are used to handle messages that do not contain a +command. This enables auto-moderation from the bot without manual intervention. The example message handler does two arbitrary things to demo functionality: 1. Checks the message content for the phrase "swear word" and, if found, times the user out for 5 minutes @@ -43,7 +45,6 @@ things to demo functionality: func main() { b, err := bolt.New() - if err != nil { panic(err) } @@ -99,8 +100,16 @@ func main() { if err != nil { panic(err) } -} + qc := make(chan os.Signal, 1) + signal.Notify(qc, os.Interrupt) + <-qc + + err = b.Stop() + if err != nil { + panic(err) + } +} ``` ## Development diff --git a/bolt.go b/bolt.go index 1d03bda..f19d36c 100644 --- a/bolt.go +++ b/bolt.go @@ -5,7 +5,6 @@ import ( "fmt" "log" "os" - "os/signal" "slices" "strings" "sync" @@ -52,6 +51,8 @@ type Bolt struct { maxRoutines int //generic message handler func msgHandlerf Payload + //command timeout when closing connection + grace time.Duration } // New creates a new bolt instance and applies any supplied options @@ -73,6 +74,7 @@ func New(opts ...Option) (*Bolt, error) { indicator: DEFAULT_INDICATOR, wg: sync.WaitGroup{}, maxRoutines: DEFAULT_MAX_GOROUTINES, + grace: time.Second * 5, } b.Identify.Intents = DEFAULT_INTENTS @@ -87,22 +89,33 @@ func New(opts ...Option) (*Bolt, error) { } // Start applies the message event handler function to the bot and opens the initial websocket connection -// with Discord. Start is a blocking call that also handles safe shutdown, on Interrupt, bolt will give -// command routines a window to finish before closing the connection. +// with Discord func (b *Bolt) Start() error { b.AddHandler(b.msgEventHandler) - err := b.Open() - if err != nil { - return fmt.Errorf("failed to open websocket connection with Discord: %e", err) + return b.Open() +} + +// AddCommands registers command handlers. Any messages that begin with the command indicator will be forwarded +// to the handler +func (b *Bolt) AddCommands(cmd ...Command) { + for _, c := range cmd { + //we can use a normal map here since writing only takes place once + b.commands[c.Trigger] = c } +} - sigChannel := make(chan os.Signal, 1) - signal.Notify(sigChannel, os.Interrupt) - <-sigChannel +// AddMessageHandler registers the generic message handler, any messages that are not commands will be forwarded +// to the handler +func (b *Bolt) AddMessageHandler(p Payload) { + b.msgHandlerf = p +} - //give handler routines a 5 second window to finish processes before closing connection - ctx, cancel := context.WithTimeout(context.Background(), time.Second*5) +// Stop waits for the alloted grace period timeout before handling safe shutdown of the +// Discord connection +func (b *Bolt) Stop() error { + ctx, cancel := context.WithTimeout(context.Background(), b.grace) defer cancel() + closeChan := make(chan struct{}) go func() { b.wg.Wait() @@ -115,25 +128,6 @@ func (b *Bolt) Start() error { case <-closeChan: } - return b.stop() -} - -// AddCommands registers command handlers. Any messages that begin with the command indicator will be forwarded -// to the handler -func (b *Bolt) AddCommands(cmd ...Command) { - for _, c := range cmd { - b.commands[c.Trigger] = c - } -} - -// AddMessageHandler registers the generic message handler, any messages that are not commands will be forwarded -// to the handler -func (b *Bolt) AddMessageHandler(p Payload) { - b.msgHandlerf = p -} - -// stop closes the websocket connection to Discord -func (b *Bolt) stop() error { return b.Close() } @@ -258,7 +252,6 @@ func (b *Bolt) handleMessage(event *Message) error { // handleCommand maps the first word of the message to the command payload, if it exists. It then forwards the message // data to the handler after checking the timeout and role restrictions. If restrictions have not been met a generic // response is sent to the message -// TODO: accept a string for timeout/role rejection messages to allow customization func (b *Bolt) handleCommand(msg *Message, lg int) error { run, ok := b.commands[msg.Words[0][lg:]] if !ok { @@ -345,8 +338,6 @@ func (b *Bolt) remainingTimeout(timeout time.Time) string { return fmt.Sprintf("%d%s", timeLeft, metric) } -// checks if the author of msg has the correct role to run the requested command - // roleCheck loops through the provided user role ID's grabs the role "name" and ensures the role is present // in the commands whitelist. If the role exists in the Command whitelist a true response is returned func (b *Bolt) roleCheck(guild string, roles []string, s *dg.Session, run Command) (bool, error) { diff --git a/option.go b/option.go index 3728b5f..0658fee 100644 --- a/option.go +++ b/option.go @@ -1,6 +1,8 @@ package bolt import ( + "time" + dg "github.com/bwmarrin/discordgo" ) @@ -48,3 +50,11 @@ func WithLogLevel(lvl LogLevel) Option { b.logLvl = lvl } } + +// WithGracePeriod gives a grace period for commands to finish prior to closing the +// websocket connection with Discord +func WithGracePeriod(t time.Duration) Option { + return func(b *Bolt) { + b.grace = t + } +}