From 770cae0006ba05c11c974ea6edf05bf70157b115 Mon Sep 17 00:00:00 2001
From: Wenda Chu <32250288+w1nda@users.noreply.github.com>
Date: Mon, 20 Jan 2025 03:54:41 +0800
Subject: [PATCH] [sonic-mgmt-docker-image] Support ptf dataplane packet poll
 with multiple ptf nn agents connection  (#21070)

When testing sonic with ptf dataplane connecting multiple ptf nn agents, some cases will fail because of packets queue in ptf were not polled thoroughly. This is a bug or missing feature in ptf: p4lang/ptf#207
as a short term quick fix, this PR will patch the ptf-py3 package and unblock our qualification process
---
 ...-poll-method-to-support-multi-ptf-nn.patch | 77 +++++++++++++++++++
 dockers/docker-sonic-mgmt/Dockerfile.j2       |  2 +
 2 files changed, 79 insertions(+)
 create mode 100644 dockers/docker-sonic-mgmt/0002-extend-dataplane-poll-method-to-support-multi-ptf-nn.patch

diff --git a/dockers/docker-sonic-mgmt/0002-extend-dataplane-poll-method-to-support-multi-ptf-nn.patch b/dockers/docker-sonic-mgmt/0002-extend-dataplane-poll-method-to-support-multi-ptf-nn.patch
new file mode 100644
index 000000000000..9d8f3f671e22
--- /dev/null
+++ b/dockers/docker-sonic-mgmt/0002-extend-dataplane-poll-method-to-support-multi-ptf-nn.patch
@@ -0,0 +1,77 @@
+From 688c4d11a00269beaf22eb6f2cccba410bfb2856 Mon Sep 17 00:00:00 2001
+From: wenda <wendachu@microsoft.com>
+Date: Fri, 6 Dec 2024 07:25:30 +0000
+Subject: [PATCH] extend dataplane poll method to support multi ptf nn agents
+ connections
+
+---
+ src/ptf/dataplane.py | 31 ++++++++++++++++++++-----------
+ 1 file changed, 20 insertions(+), 11 deletions(-)
+
+diff --git a/src/ptf/dataplane.py b/src/ptf/dataplane.py
+index a1c1b3f..009ac2f 100644
+--- a/src/ptf/dataplane.py
++++ b/src/ptf/dataplane.py
+@@ -738,40 +738,49 @@ class DataPlane(Thread):
+             )
+         return bytes
+
+-    def oldest_port_number(self, device):
++    def get_oldest_tuple(self, device):
+         """
+-        Returns the port number with the oldest packet,
++        Returns the device number and port number with the oldest packet,
+         or None if no packets are queued.
++        When device is specified, only returns the oldest packet from that device.
+         """
+-        min_port_number = None
++        min_device_number, min_port_number = None, None
+         min_time = float("inf")
+         for port_id, queue in list(self.packet_queues.items()):
+-            if port_id[0] != device:
++            if device and port_id[0] != device:
+                 continue
+             if queue and queue[0][1] < min_time:
+                 min_time = queue[0][1]
++                min_device_number = port_id[0]
+                 min_port_number = port_id[1]
+-        return min_port_number
++        return min_device_number, min_port_number
+
+     # Dequeues and yields packets in the order they were received.
+     # Yields (port, packet, received time).
+     # If port is not specified yields packets from all ports.
++    # If port and device are both not specified yields packets from all devices and all ports
+     def packets(self, device, port=None):
+         while True:
+-            if port is None:
+-                rcv_port = self.oldest_port_number(device)
+-            else:
+-                rcv_port = port
++            rcv_device, rcv_port = device, port
++            if device is None and port is None:
++                rcv_device, rcv_port = self.get_oldest_tuple(None)
++            elif port is None:
++                _, rcv_port = self.get_oldest_tuple(device)
++            elif device is None:
++                self.logger.error(
++                    "ambiguous tuple given. device is None, while port is %s" % (port)
++                )
++                break
+
+             if rcv_port == None:
+                 self.logger.debug("Out of packets on all ports")
+                 break
+
+-            queue = self.packet_queues[(device, rcv_port)]
++            queue = self.packet_queues[(rcv_device, rcv_port)]
+
+             if len(queue) == 0:
+                 self.logger.debug(
+-                    "Out of packets on device %d, port %d", device, rcv_port
++                    "Out of packets on device %d, port %d", rcv_device, rcv_port
+                 )
+                 break
+
+-- 
+2.47.0
\ No newline at end of file
diff --git a/dockers/docker-sonic-mgmt/Dockerfile.j2 b/dockers/docker-sonic-mgmt/Dockerfile.j2
index ecd4e6a54fab..64badcb8eb83 100755
--- a/dockers/docker-sonic-mgmt/Dockerfile.j2
+++ b/dockers/docker-sonic-mgmt/Dockerfile.j2
@@ -271,11 +271,13 @@ USER root
 WORKDIR /azp
 COPY start.sh \
      0001-Fix-getattr-AttributeError-in-multi-thread-scenario.patch \
+     0002-extend-dataplane-poll-method-to-support-multi-ptf-nn.patch \
      ./
 RUN chmod +x start.sh \
     && ln -sf /usr/bin/python3 /usr/bin/python \
     && ln -sf `which pip3` /usr/bin/pip \
     && ln -sf `which pip3` /usr/local/sbin/pip \
+    && patch -u -b /usr/local/lib/python3.8/dist-packages/ptf/dataplane.py  -i /azp/0002-extend-dataplane-poll-method-to-support-multi-ptf-nn.patch \
     && patch -u -b /usr/local/lib/python3.8/dist-packages/ansible/plugins/loader.py -i /azp/0001-Fix-getattr-AttributeError-in-multi-thread-scenario.patch
 
 USER $user