From c053c1e8c45598410926e2ac9aec560520b13be2 Mon Sep 17 00:00:00 2001
From: MoofMonkey <11695747+MoofMonkey@users.noreply.github.com>
Date: Tue, 4 Oct 2022 03:16:12 +0300
Subject: [PATCH] Fix flaky TestManagerCtxDoneCloseSessions

---
 datagramsession/manager.go | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/datagramsession/manager.go b/datagramsession/manager.go
index a461df4c..a315909c 100644
--- a/datagramsession/manager.go
+++ b/datagramsession/manager.go
@@ -2,6 +2,7 @@ package datagramsession
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"time"
@@ -87,9 +88,10 @@ func (m *manager) shutdownSessions(err error) {
 	}
 	closeSessionErr := &errClosedSession{
 		message: err.Error(),
-		// Usually connection with remote has been closed, so set this to true to skip unregistering from remote
-		byRemote: true,
 	}
+	// Usually connection with remote has been closed, so set this to true to skip unregistering from remote
+	// context.Canceled is an exception because that means session is being closed by our side
+	closeSessionErr.byRemote = !errors.Is(err, context.Canceled)
 	for _, s := range m.sessions {
 		s.close(closeSessionErr)
 	}