parent
e9678cee2f
commit
d423a1ca24
25
conn.go
25
conn.go
@ -143,6 +143,10 @@ func newConn(srv *Server, ic ircConn, options *connOptions) *conn {
|
||||
|
||||
rl := rate.NewLimiter(rate.Every(options.RateLimitDelay), options.RateLimitBurst)
|
||||
for msg := range outgoing {
|
||||
if msg == nil {
|
||||
break
|
||||
}
|
||||
|
||||
if err := rl.Wait(ctx); err != nil {
|
||||
break
|
||||
}
|
||||
@ -224,6 +228,27 @@ func (c *conn) SendMessage(ctx context.Context, msg *irc.Message) {
|
||||
}
|
||||
}
|
||||
|
||||
// Shutdown gracefully closes the connection, flushing any pending message.
|
||||
func (c *conn) Shutdown(ctx context.Context) {
|
||||
c.lock.Lock()
|
||||
defer c.lock.Unlock()
|
||||
|
||||
if c.closed {
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case c.outgoing <- nil:
|
||||
// Success
|
||||
case <-ctx.Done():
|
||||
c.logger.Printf("failed to shutdown connection: %v", ctx.Err())
|
||||
// Forcibly close the connection
|
||||
if err := c.Close(); err != nil && !errors.Is(err, net.ErrClosed) {
|
||||
c.logger.Printf("failed to close connection: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *conn) RemoteAddr() net.Addr {
|
||||
return c.conn.RemoteAddr()
|
||||
}
|
||||
|
@ -622,7 +622,8 @@ func (dc *downstreamConn) handleMessage(ctx context.Context, msg *irc.Message) e
|
||||
|
||||
switch msg.Command {
|
||||
case "QUIT":
|
||||
return dc.Close()
|
||||
dc.conn.Shutdown(ctx)
|
||||
return nil // TODO: stop handling commands
|
||||
default:
|
||||
if dc.registered {
|
||||
return dc.handleMessageRegistered(ctx, msg)
|
||||
@ -1698,7 +1699,7 @@ func (dc *downstreamConn) runUntilRegistered() error {
|
||||
Command: "ERROR",
|
||||
Params: []string{"Connection registration timed out"},
|
||||
})
|
||||
dc.Close()
|
||||
dc.Shutdown(ctx)
|
||||
}
|
||||
}()
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user