docs: def drain(self, cluster, nodes, max_size):
          """
          Drain all the data for the given nodes and collate them into a list of
          batches that will fit within the specified size on a per-node basis.
          This method attempts to avoid choosing the same topic-node repeatedly.
  
          Arguments:
              cluster (ClusterMetadata): The current cluster metadata
              nodes (list): list of node_ids to drain
              max_size (int): maximum number of bytes to drain
  
          Returns:
              dict: {node_id: list of ProducerBatch} with total size less than the
                  requested max_size.
          """
          if not nodes:
              return {}
  
          now = time.time()
          batches = {}
          for node_id in nodes:
              size = 0
              partitions = list(cluster.partitions_for_broker(node_id))
              ready = []
              # to make starvation less likely this loop doesn't start at 0
              self._drain_index %= len(partitions)
              start = self._drain_index
              while True:
                  tp = partitions[self._drain_index]
                  if tp in self._batches and tp not in self.muted:
                      with self._tp_locks[tp]:
                          dq = self._batches[tp]
                          if dq:
                              first = dq[0]
                              backoff = (
                                  bool(first.attempts > 0) and
                                  bool(first.last_attempt +
                                       self.config['retry_backoff_ms'] / 1000.0
                                       > now)
                              )
                              # Only drain the batch if it is not during backoff
                              if not backoff:
                                  if (size + first.records.size_in_bytes() > max_size
                                      and len(ready) > 0):
                                      # there is a rare case that a single batch
                                      # size is larger than the request size due
                                      # to compression; in this case we will
                                      # still eventually send this batch in a
                                      # single request
                                      break
                                  else:
                                      batch = dq.popleft()
                                      batch.records.close()
                                      size += batch.records.size_in_bytes()
                                      ready.append(batch)
                                      batch.drained = now
  
                  self._drain_index += 1
                  self._drain_index %= len(partitions)
                  if start == self._drain_index:
                      break
  
              batches[node_id] = ready
          return batches