'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 RecordBatch} with total size less than the
requested max_size.'

def drain(self, cluster, nodes, 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 = []
        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)))
                        if (not backoff):
                            if (((size + first.records.size_in_bytes()) > max_size) and (len(ready) > 0)):
                                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
