Introduce conn.NewContext

This function wraps a parent context, and returns a new context
cancelled when the connection is closed. This will make it so
operations started from downstreamConn.handleMessage will be
cancelled when the connection is closed.
This commit is contained in:
Simon Ser 2021-11-17 13:13:55 +01:00
parent eb47f0dcd7
commit 6143e6f12d
2 changed files with 32 additions and 1 deletions

28
conn.go
View File

@ -172,6 +172,7 @@ type conn struct {
lock sync.Mutex
outgoing chan<- *irc.Message
closed bool
closedCh chan struct{}
}
func newConn(srv *Server, ic ircConn, options *connOptions) *conn {
@ -181,6 +182,7 @@ func newConn(srv *Server, ic ircConn, options *connOptions) *conn {
srv: srv,
outgoing: outgoing,
logger: options.Logger,
closedCh: make(chan struct{}),
}
go func() {
@ -237,6 +239,7 @@ func (c *conn) Close() error {
err := c.conn.Close()
c.closed = true
close(c.outgoing)
close(c.closedCh)
return err
}
@ -277,3 +280,28 @@ func (c *conn) RemoteAddr() net.Addr {
func (c *conn) LocalAddr() net.Addr {
return c.conn.LocalAddr()
}
// NewContext returns a copy of the parent context with a new Done channel. The
// returned context's Done channel is closed when the connection is closed,
// when the returned cancel function is called, or when the parent context's
// Done channel is closed, whichever happens first.
//
// Canceling this context releases resources associated with it, so code should
// call cancel as soon as the operations running in this Context complete.
func (c *conn) NewContext(parent context.Context) (context.Context, context.CancelFunc) {
ctx, cancel := context.WithCancel(parent)
go func() {
defer cancel()
select {
case <-ctx.Done():
// The parent context has been cancelled, or the caller has called
// cancel()
case <-c.closedCh:
// The connection has been closed
}
}()
return ctx, cancel
}

View File

@ -615,7 +615,10 @@ func (dc *downstreamConn) marshalMessage(msg *irc.Message, net *network) *irc.Me
}
func (dc *downstreamConn) handleMessage(msg *irc.Message) error {
ctx, cancel := context.WithTimeout(context.TODO(), handleDownstreamMessageTimeout)
ctx, cancel := dc.conn.NewContext(context.TODO())
defer cancel()
ctx, cancel = context.WithTimeout(ctx, handleDownstreamMessageTimeout)
defer cancel()
switch msg.Command {