Replace time.After with time.NewTimer

Improper use of time.After can lead to memory leaks if the timer never
gets a chance to fire.

Signed-off-by: Gabriel Adrian Samfira <gsamfira@cloudbasesolutions.com>
This commit is contained in:
Gabriel Adrian Samfira 2024-07-05 10:48:27 +00:00
parent 49f1b7aa14
commit 2554f70b89
6 changed files with 32 additions and 16 deletions

View file

@ -134,9 +134,11 @@ func (c *Client) Write(msg []byte) (int, error) {
tmp := make([]byte, len(msg))
copy(tmp, msg)
timer := time.NewTimer(5 * time.Second)
defer timer.Stop()
select {
case <-time.After(5 * time.Second):
case <-timer.C:
return 0, fmt.Errorf("timed out sending message to client")
case c.send <- tmp:
}
@ -193,11 +195,6 @@ func (c *Client) writeMessage(messageType int, message []byte) error {
// clientWriter
func (c *Client) clientWriter() {
ticker := time.NewTicker(pingPeriod)
defer func() {
c.Stop()
ticker.Stop()
}()
// Set up expiration timer.
// NOTE: if a token is created without an expiration date
// this will be set to nil, which will close the loop bellow
@ -208,6 +205,13 @@ func (c *Client) clientWriter() {
if expires != nil {
authExpires = *expires
}
authTimer := time.NewTimer(time.Until(authExpires))
ticker := time.NewTicker(pingPeriod)
defer func() {
c.Stop()
ticker.Stop()
authTimer.Stop()
}()
for {
select {
case message, ok := <-c.send:
@ -236,7 +240,7 @@ func (c *Client) clientWriter() {
}
case <-c.ctx.Done():
return
case <-time.After(time.Until(authExpires)):
case <-authTimer.C:
// Auth has expired
slog.DebugContext(c.ctx, "auth expired, closing connection")
return

View file

@ -107,9 +107,10 @@ func (h *Hub) Unregister(client *Client) error {
func (h *Hub) Write(msg []byte) (int, error) {
tmp := make([]byte, len(msg))
copy(tmp, msg)
timer := time.NewTimer(5 * time.Second)
defer timer.Stop()
select {
case <-time.After(5 * time.Second):
case <-timer.C:
return 0, fmt.Errorf("timed out sending message to client")
case h.broadcast <- tmp:
}
@ -134,9 +135,11 @@ func (h *Hub) Stop() error {
}
func (h *Hub) Wait() error {
timer := time.NewTimer(60 * time.Second)
defer timer.Stop()
select {
case <-h.closed:
case <-time.After(60 * time.Second):
case <-timer.C:
return fmt.Errorf("timed out waiting for hub stop")
}
return nil