From 6b25d26ea87c1d279c8f7347f7b571fabc4fa82a Mon Sep 17 00:00:00 2001
From: Aleksey Myasnikov <asmyasnikov@ydb.tech>
Date: Mon, 3 Feb 2025 18:02:57 +0300
Subject: [PATCH] wrapped closed context on closing sessions

---
 internal/query/session_core.go |  4 ++++
 internal/table/session.go      | 11 ++++++++++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/internal/query/session_core.go b/internal/query/session_core.go
index e0f1ca768..26c4e1cff 100644
--- a/internal/query/session_core.go
+++ b/internal/query/session_core.go
@@ -234,6 +234,10 @@ func (core *sessionCore) deleteSession(ctx context.Context) (finalErr error) {
 		onDone(finalErr)
 	}()
 
+	if err := ctx.Err(); err != nil {
+		return xerrors.WithStackTrace(err)
+	}
+
 	if d := core.deleteTimeout; d > 0 {
 		var cancel context.CancelFunc
 		ctx, cancel = xcontext.WithTimeout(ctx, d)
diff --git a/internal/table/session.go b/internal/table/session.go
index 063ddf6ea..d0727d054 100644
--- a/internal/table/session.go
+++ b/internal/table/session.go
@@ -357,6 +357,10 @@ func newTableSession( //nolint:funlen
 		status: table.SessionReady,
 		onClose: []func(s *Session) error{
 			func(s *Session) error {
+				if err := ctx.Err(); err != nil {
+					return xerrors.WithStackTrace(err)
+				}
+
 				_, err = s.client.DeleteSession(ctx,
 					&Ydb_Table.DeleteSessionRequest{
 						SessionId: s.id,
@@ -481,13 +485,18 @@ func (s *Session) Close(ctx context.Context) (err error) {
 		s.SetStatus(table.SessionClosed)
 	}()
 
+	errs := make([]error, 0, len(s.onClose))
 	for _, onClose := range s.onClose {
 		err := onClose(s)
 		if err != nil {
-			return xerrors.WithStackTrace(err)
+			errs = append(errs, err)
 		}
 	}
 
+	if len(errs) > 0 {
+		return xerrors.WithStackTrace(xerrors.Join(errs...))
+	}
+
 	return nil
 }