From f61df169c714fc5ce4b5bce560a52faea7d1a887 Mon Sep 17 00:00:00 2001
From: Matt Broadstone <mbroadst@mongodb.com>
Date: Wed, 17 Jul 2019 07:47:26 -0400
Subject: [PATCH] feat(sessions): track dirty state of sessions, drop after use

If a network error occurs during an operation using a session, that
session must now be marked dirty and discarded after the operation
completes (allowing for retryability).

NODE-1977
---
 lib/core/sdam/server.js | 24 ++++++++++++++++++++++++
 lib/core/sessions.js    | 10 ++++++++--
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/lib/core/sdam/server.js b/lib/core/sdam/server.js
index 0a1dbb5a83..2feeea7c0e 100644
--- a/lib/core/sdam/server.js
+++ b/lib/core/sdam/server.js
@@ -240,6 +240,12 @@ class Server extends EventEmitter {
     }
 
     wireProtocol.command(this, ns, cmd, options, (err, result) => {
+      if (err && err instanceof MongoNetworkError) {
+        if (options.session) {
+          options.session.serverSession.isDirty = true;
+        }
+      }
+
       if (err && isSDAMUnrecoverableError(err)) {
         this.emit('error', err);
       }
@@ -258,6 +264,12 @@ class Server extends EventEmitter {
    */
   query(ns, cmd, cursorState, options, callback) {
     wireProtocol.query(this, ns, cmd, cursorState, options, (err, result) => {
+      if (err && err instanceof MongoNetworkError) {
+        if (options.session) {
+          options.session.serverSession.isDirty = true;
+        }
+      }
+
       if (err && isSDAMUnrecoverableError(err)) {
         this.emit('error', err);
       }
@@ -276,6 +288,12 @@ class Server extends EventEmitter {
    */
   getMore(ns, cursorState, batchSize, options, callback) {
     wireProtocol.getMore(this, ns, cursorState, batchSize, options, (err, result) => {
+      if (err && err instanceof MongoNetworkError) {
+        if (options.session) {
+          options.session.serverSession.isDirty = true;
+        }
+      }
+
       if (err && isSDAMUnrecoverableError(err)) {
         this.emit('error', err);
       }
@@ -406,6 +424,12 @@ function executeWriteOperation(args, options, callback) {
   }
 
   return wireProtocol[op](server, ns, ops, options, (err, result) => {
+    if (err && err instanceof MongoNetworkError) {
+      if (options.session) {
+        options.session.serverSession.isDirty = true;
+      }
+    }
+
     if (err && isSDAMUnrecoverableError(err)) {
       server.emit('error', err);
     }
diff --git a/lib/core/sessions.js b/lib/core/sessions.js
index 446f9933e2..f77e2188b3 100644
--- a/lib/core/sessions.js
+++ b/lib/core/sessions.js
@@ -527,6 +527,7 @@ class ServerSession {
     this.id = { id: new Binary(uuidV4(), Binary.SUBTYPE_UUID) };
     this.lastUse = Date.now();
     this.txnNumber = 0;
+    this.isDirty = false;
   }
 
   /**
@@ -603,8 +604,8 @@ class ServerSessionPool {
   release(session) {
     const sessionTimeoutMinutes = this.topology.logicalSessionTimeoutMinutes;
     while (this.sessions.length) {
-      const session = this.sessions[this.sessions.length - 1];
-      if (session.hasTimedOut(sessionTimeoutMinutes)) {
+      const pooledSession = this.sessions[this.sessions.length - 1];
+      if (pooledSession.hasTimedOut(sessionTimeoutMinutes)) {
         this.sessions.pop();
       } else {
         break;
@@ -612,6 +613,11 @@ class ServerSessionPool {
     }
 
     if (!session.hasTimedOut(sessionTimeoutMinutes)) {
+      if (session.isDirty) {
+        return;
+      }
+
+      // otherwise, readd this session to the session pool
       this.sessions.unshift(session);
     }
   }