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.

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