/*
 Licensed to the Apache Software Foundation (ASF) under one
 or more contributor license agreements.  See the NOTICE file
 distributed with this work for additional information
 regarding copyright ownership.  The ASF licenses this file
 to you under the Apache License, Version 2.0 (the
 "License"); you may not use this file except in compliance
 with the License.  You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 */

package org.apache.flink.runtime.source.coordinator;

import org.apache.flink.api.connector.source.ReaderInfo;
import org.apache.flink.core.io.SimpleVersionedSerializer;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.nio.BufferUnderflowException;
import java.util.HashMap;
import java.util.LinkedHashSet;
import java.util.Map;

/**
 * A serialization util class for the {@link SourceCoordinator}.
 */
public class SourceCoordinatorSerdeUtils {
	/** The current source coordinator serde version. */
	private static final int CURRENT_VERSION = 0;

	/** Private constructor for utility class. */
	private SourceCoordinatorSerdeUtils() {}

	/** Write the current serde version. */
	static void writeCoordinatorSerdeVersion(DataOutputStream out) throws IOException {
		out.writeInt(CURRENT_VERSION);
	}

	/** Read and verify the serde version. */
	static void readAndVerifyCoordinatorSerdeVersion(DataInputStream in) throws IOException {
		int version = in.readInt();
		if (version > CURRENT_VERSION) {
			throw new IOException("Unsupported source coordinator serde version " + version);
		}
	}

	/**
	 * Get serialized size of the registered readers map.
	 *
	 * <p>The binary format is following:
	 * 4 Bytes - num entries.
	 * N Bytes - entries
	 * 		4 Bytes - subtask id
	 * 		N Bytes - reader info, see {@link #writeReaderInfo(ReaderInfo, DataOutputStream)}.
	 */
	static void writeRegisteredReaders(Map<Integer, ReaderInfo> registeredReaders, DataOutputStream out) throws IOException {
		out.writeInt(registeredReaders.size());
		for (ReaderInfo info : registeredReaders.values()) {
			writeReaderInfo(info, out);
		}
	}

	static Map<Integer, ReaderInfo> readRegisteredReaders(DataInputStream in) throws IOException {
		int numReaders = in.readInt();
		Map<Integer, ReaderInfo> registeredReaders = new HashMap<>();
		for (int i = 0; i < numReaders; i++) {
			ReaderInfo info = readReaderInfo(in);
			registeredReaders.put(info.getSubtaskId(), info);
		}
		return registeredReaders;
	}

	/**
	 * Serialize the assignment by checkpoint ids.
	 */
	static <SplitT> void writeAssignmentsByCheckpointId(
			Map<Long, Map<Integer, LinkedHashSet<SplitT>>> assignmentByCheckpointIds,
			SimpleVersionedSerializer<SplitT> splitSerializer,
			DataOutputStream out) throws IOException {
		// SplitSerializer version.
		out.writeInt(splitSerializer.getVersion());
		// Num checkpoints.
		out.writeInt(assignmentByCheckpointIds.size());
		for (Map.Entry<Long, Map<Integer, LinkedHashSet<SplitT>>> assignments : assignmentByCheckpointIds.entrySet()) {
			long checkpointId = assignments.getKey();
			out.writeLong(checkpointId);

			int numSubtasks = assignments.getValue().size();
			out.writeInt(numSubtasks);
			for (Map.Entry<Integer, LinkedHashSet<SplitT>> assignment : assignments.getValue().entrySet()) {
				int subtaskId = assignment.getKey();
				out.writeInt(subtaskId);

				int numAssignedSplits = assignment.getValue().size();
				out.writeInt(numAssignedSplits);
				for (SplitT split : assignment.getValue()) {
					byte[] serializedSplit = splitSerializer.serialize(split);
					out.writeInt(serializedSplit.length);
					out.write(serializedSplit);
				}
			}
		}
	}

	/**
	 * Deserialize the assignment by checkpoint ids.
	 */
	static <SplitT> Map<Long, Map<Integer, LinkedHashSet<SplitT>>> readAssignmentsByCheckpointId(
			DataInputStream in,
			SimpleVersionedSerializer<SplitT> splitSerializer) throws IOException {
		int splitSerializerVersion = in.readInt();
		int numCheckpoints = in.readInt();
		Map<Long, Map<Integer, LinkedHashSet<SplitT>>> assignmentsByCheckpointIds = new HashMap<>(numCheckpoints);
		for (int i = 0; i < numCheckpoints; i++) {
			long checkpointId = in.readLong();
			int numSubtasks = in.readInt();
			Map<Integer, LinkedHashSet<SplitT>> assignments = new HashMap<>();
			assignmentsByCheckpointIds.put(checkpointId, assignments);
			for (int j = 0; j < numSubtasks; j++) {
				int subtaskId = in.readInt();
				int numAssignedSplits = in.readInt();
				LinkedHashSet<SplitT> splits = new LinkedHashSet<>(numAssignedSplits);
				assignments.put(subtaskId, splits);
				for (int k = 0; k < numAssignedSplits; k++) {
					int serializedSplitSize = in.readInt();
					byte[] serializedSplit = readBytes(in, serializedSplitSize);
					SplitT split = splitSerializer.deserialize(splitSerializerVersion, serializedSplit);
					splits.add(split);
				}
			}
		}
		return assignmentsByCheckpointIds;
	}

	static byte[] readBytes(DataInputStream in, int size) throws IOException {
		byte[] bytes = new byte[size];
		int off = 0;
		// For ByteArrayInputStream the read should succeed with one shot.
		while (off < size) {
			int read = in.read(bytes, off, size - off);
			if (read < 0) {
				throw new BufferUnderflowException();
			}
			off += read;
		}
		return bytes;
	}

	// ----- private helper methods -----

	/**
	 * Serialize {@link ReaderInfo}.
	 *
	 * <p>The binary format is following:
	 * 4 Bytes - subtask id
	 * N Bytes - location string
	 *
	 * @param readerInfo the given reader information to serialize.
	 */
	private static void writeReaderInfo(ReaderInfo readerInfo, DataOutputStream out) throws IOException {
		out.writeInt(readerInfo.getSubtaskId());
		out.writeUTF(readerInfo.getLocation());
	}

	private static ReaderInfo readReaderInfo(DataInputStream in) throws IOException {
		int subtaskId = in.readInt();
		String location = in.readUTF();
		return new ReaderInfo(subtaskId, location);
	}

}

@@ -1040,50 +1040,8 @@
 on;%0A
-import java.nio.BufferUnderflowException;%0A
 impo
@@ -5337,243 +5337,26 @@
 %09%09in
-t off = 0;%0A%09%09// For ByteArrayInputStream the read should succeed with one shot.%0A%09%09while (off %3C size) %7B%0A%09%09%09int read = in.read(bytes, off, size - off);%0A%09%09%09if (read %3C 0) %7B%0A%09%09%09%09throw new BufferUnderflowException();%0A%09%09%09%7D%0A%09%09%09off += read;%0A%09%09%7D
+.readFully(bytes);
 %0A%09%09r

Replace manual byte[] reading with DataInput.readFully()