Fixed issue of incorrect splitting of response buffer into multiple packets for PLP Data (#3830)

Currently, in Babelfish for large response of nvarchar(max) type while splitting the response buffer into multiple packets the last PLP chunk is getting splitting directly and the next packet starts with rest of the part of PLP chunk. Instead if the last PLP chunk size cannot be accommodated into available buffer, the PLP chunk size should have been reduced to available buffer and next packet should start with new PLP chunk. This PR will fix this issue.

Task: BABEL-5875

Signed-off-by: Rohit Bhagat <rohitbgt@amazon.com>

Author: rohit01010

contrib/babelfishpg_tds/src/backend/tds/tdscomm.c

@@ -58,6 +58,31 @@ static void SocketSetNonblocking(bool nonblocking);
 static int	InternalFlush(bool);
 static void TdsConsumedBytes(int bytes);
 
+/* --------------------------------
+ *	GetAvailableBufferSize - returns the available buffer size
+ * --------------------------------
+ */
+size_t
+GetAvailableBufferSize()
+{
+	return TdsBufferSize - TdsSendCur;
+}
+
+/* --------------------------------
+ *	FlushBuffer - Flushes existing buffer with socket set as non-blocking
+ *
+ *	returns 0 if OK, EOF if trouble
+ * --------------------------------
+ */
+int
+FlushBuffer()
+{
+	SocketSetNonblocking(false);
+	if (InternalFlush(false))
+		return EOF;
+	return 0;
+}
+
 /* Inline functions */
 
 /* --------------------------------

contrib/babelfishpg_tds/src/backend/tds/tdstypeio.c

@@ -2810,18 +2810,51 @@ TdsSendPlpDataHelper(char *data, int len)
 	uint32_t	plpTerminator = PLP_TERMINATOR;
 	uint64_t	tempOffset = 0;
 	uint32_t	plpChunckLen = PLP_CHUNCK_LEN;
+	size_t		availableBufferSpace = 0;
 
+	TdsErrorContext->err_text = "TdsSendPlpDataHelper - write PLP Data to the buffer";
+	TDS_DEBUG(TDS_DEBUG2, "PLP Data length = %d", len);
 	if ((rc = TdsPutInt64LE(len)) == 0)
 	{
 		while (true)
 		{
+			plpChunckLen = PLP_CHUNCK_LEN;
+
 			if (plpChunckLen > (len - tempOffset))
 				plpChunckLen = (len - tempOffset);
 
 			/* Either data is "0" or no more data to send */
 			if (plpChunckLen == 0)
 				break;
+			
+			availableBufferSpace = GetAvailableBufferSize();
+			/*
+			 * If PLP CHUNK block cannot be stored in available buffer, 
+			 * then set PLP CHUNK LEN to size of available buffer.
+			 * 
+			 * Here plpChunckLen + sizeof(uint32_t) represents the size of chunk data 
+			 * and its length(i.e. size of uint32_t). This check will make sure that 
+			 * each PLP CHUNK and its length will remain in same packet.
+			 */
+			if ((plpChunckLen + sizeof(uint32_t)) > availableBufferSpace)
+			{
+				/*
+				 * If available buffer size cannot accomodate smallest chunk i.e. of size 1 byte
+				 * then flush the buffer and try again, else set PLP CHUNK LEN to size of available buffer.
+				 */
+				if (availableBufferSpace < (1 + sizeof(uint32_t)))
+				{
+					rc = FlushBuffer();
+					if (rc != 0)
+						return rc;
+
+					continue;
+				}
+				else
+					plpChunckLen = availableBufferSpace - sizeof(uint32_t);
+			}
 
+			TDS_DEBUG(TDS_DEBUG2, "PLP CHUNK length = %u, Available Buffer Size = %zu", plpChunckLen, availableBufferSpace);
 			/* need testing for "0" len */
 			if ((rc = TdsPutUInt32LE(plpChunckLen)) == 0)
 			{

contrib/babelfishpg_tds/src/include/tds_int.h

@@ -306,6 +306,8 @@ extern uint32_t GetClientTDSVersion(void);
 extern char *get_tds_login_domainname(void);
 extern void TdsSetDbContext(void);
 extern void TdsResetLoginFlags(void);
+extern size_t GetAvailableBufferSize(void);
+extern int FlushBuffer(void);
 
 /* Functions in backend/tds/tdsprotocol.c */
 extern int	TdsSocketBackend(void);