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:
parent
eb47f0dcd7
commit
6143e6f12d
28
conn.go
28
conn.go
@ -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
|
||||
}
|
||||
|
@ -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 {
|
||||
|
Loading…
Reference in New Issue
Block a user