Minor changes made to adapt to new spec points and failures found

Controller.java

@@ -192,12 +192,7 @@ public class Controller {
 	protected class RebalanceThread implements Runnable {
 		public void run() {
 			while(true) {
-				if(rebalanceLock.waitToRebalance()) {
-					//Another dstore joined, it requested a rebalance
-					try {runRebalance();} catch(Exception e) {e.printStackTrace();}
-				}
-				else {
-					//Timeout occured, i.e. rebalancePeriod has passed since the last rebalance
+				rebalanceLock.waitToRebalance();
 				try {
 					if(dstores.size() >= rFactor) {
 						runRebalance();
@@ -206,7 +201,6 @@ public class Controller {
 				catch(Exception e) {e.printStackTrace();}
 			}
 		}
-		}
 		
 		protected void runRebalance() {
 			synchronized(rebalanceLock) {
@@ -312,12 +306,13 @@ public class Controller {
 			String message = Protocol.STORE_TO_TOKEN;
 			for(Integer thisStore : storesToStore) {
 				message = message + " " + thisStore.intValue();
+				IndexEntry entryf = entry;
 				new Thread(() -> {
 					try {
-						String[] receivedMessage = dstores.get(thisStore).receive(Protocol.STORE_ACK_TOKEN).split(" ");
+						String[] receivedMessage = dstores.get(thisStore).receive(Protocol.STORE_ACK_TOKEN + " " + filename).split(" ");
 						if(receivedMessage[0].equals(Protocol.STORE_ACK_TOKEN)) {
 							try {
-								storeAck(thisStore, receivedMessage[1], latch);
+								storeAck(thisStore, entryf, latch);
 							}
 							catch(Exception e) {
 								//Log error
@@ -334,6 +329,9 @@ public class Controller {
 						e.printStackTrace();
 						removeDstore(e);
 					}
+					catch(DeadStoreException e) {
+						System.err.println("Store for " + filename + " failed due to dead dstore");
+					}
 				}).start();
 			}
 			out.println(message);
@@ -365,14 +363,8 @@ public class Controller {
 		}
 	}
 	
-	void storeAck(Integer port, String filename, CountDownLatch latch) throws Exception {
-		if(!index.containsKey(filename)) {
-			//Throw logging exception
-			throw new Exception("Index does not contain " + filename);
-		}
-			
-		IndexEntry thisEntry = index.get(filename);
-		thisEntry.addStoredBy(port);
+	void storeAck(Integer port, IndexEntry entry, CountDownLatch latch) throws Exception {
+		entry.addStoredBy(port);
 		latch.countDown();
 	}
 	
@@ -409,9 +401,13 @@ public class Controller {
 	void sendLoadFrom(Socket client, String filename) {
 		try {
 			PrintWriter out = new PrintWriter(client.getOutputStream());
+			String message;
+			if(!index.containsKey(filename) || index.get(filename).getStatus() != IndexEntry.Status.STORE_COMPLETE) {
+				message = Protocol.ERROR_FILE_DOES_NOT_EXIST_TOKEN;
+			}
+			else {
 				Reloader storedBy = loadRequests.get(client);
 				System.out.println("Load requested for file " + filename + ", there are " + storedBy.size() + " dstores to select from");
-			String message;
 				if(storedBy.isEmpty()) {
 					message = Protocol.ERROR_LOAD_TOKEN;
 				}
@@ -420,6 +416,7 @@ public class Controller {
 					storedBy.remove(thisStore);
 					message = Protocol.LOAD_FROM_TOKEN + thisStore + " " + storedBy.filesize;
 				}
+			}
 			out.println(message);
 			out.flush();
 			messageSent(client, message);
@@ -464,7 +461,7 @@ public class Controller {
 				Integer dstore = it.next();
 				new Thread(() -> {
 					try {
-						String[] message = dstores.get(dstore).sendAndReceive(Protocol.REMOVE_TOKEN + " " + filename, Protocol.REMOVE_ACK_TOKEN, Protocol.ERROR_FILE_DOES_NOT_EXIST_TOKEN).split(" ");
+						String[] message = dstores.get(dstore).sendAndReceive(Protocol.REMOVE_TOKEN + " " + filename, Protocol.REMOVE_ACK_TOKEN + " " + filename, Protocol.ERROR_FILE_DOES_NOT_EXIST_TOKEN).split(" ");
 						if((message[0].equals(Protocol.REMOVE_ACK_TOKEN) && message[1].equals(filename)) || message[0].equals(Protocol.ERROR_FILE_DOES_NOT_EXIST_TOKEN)) {
 							entry.removeStoredBy(dstore.intValue());
 							latch.countDown();
@@ -478,15 +475,14 @@ public class Controller {
 						e.printStackTrace();
 						removeDstore(e);
 					}
+					catch(DeadStoreException e) {
+						System.err.println("Remove for " + filename + " failed due to dead dstore");
+					}
 				}).start();
 			}
 			
 			//Wait for REMOVE_ACKs from all Dstores which were sent the REMOVE message
-			if(!latch.await(timeout, TimeUnit.MILLISECONDS)) {
-				//Log error
-				System.err.println("Not all REMOVE_ACKs have been received");
-			}
-			
+			if(latch.await(timeout, TimeUnit.MILLISECONDS)) {
 				//Update index to "remove complete"
 				entry.setStatus(IndexEntry.Status.REMOVE_COMPLETE);
 				synchronized(index) {
@@ -499,6 +495,11 @@ public class Controller {
 				clientOut.flush();
 				messageSent(client, Protocol.REMOVE_COMPLETE_TOKEN);
 			}
+			else {
+				//Log error
+				System.err.println("Not all REMOVE_ACKs have been received");
+			}
+		}
 		catch(IOException e) {
 			e.printStackTrace();
 		}
@@ -549,6 +550,7 @@ public class Controller {
 						removeDstore(e);
 						listLatch.countDown();
 					}
+					catch(DeadStoreException e) {}
 				});
 				thisThread.start();
 				activeThreads.add(thisThread);
@@ -572,7 +574,7 @@ public class Controller {
 			}
 			catch(Exception e) {e.printStackTrace();}
 			
-			if(dstoreFiles.isEmpty()) throw new Exception("All dstores have been disconnected!");
+			if(dstoreFiles.size() < rFactor) throw new Exception("Less than R dstores connected; connections may be faulty or timeout may be too strict");
 			
 			Map<Integer,List<String>> newAlloc = allocate(dstoreFiles);
 			Map<Integer,String> sendIndex = composeRebalanceMessages(dstoreFiles, newAlloc);
@@ -644,7 +646,14 @@ public class Controller {
 		for(Integer i : oldDstoreFiles.keySet()) {
 			List<String> files = new ArrayList<String>();
 			for(String s : oldDstoreFiles.get(i)) {
-				if(index.containsKey(s)) files.add(s);
+				if(index.containsKey(s)) {
+					if(index.get(s).getStatus() == IndexEntry.Status.STORE_COMPLETE) {
+						files.add(s);
+					}
+					else {
+						index.remove(s);
+					}
+				}
 				if(!availableFiles.contains(s)) availableFiles.add(s);
 			}
 			dstoreFiles.put(i, files);

DeadStoreException.java

+import java.lang.Throwable;
+import java.net.Socket;
+
+public class DeadStoreException extends Exception {
+	DstoreConnection connection;
+	
+	public DeadStoreException(DstoreConnection connection) {
+		super("Dstore at port " + connection.getPort() + " is unavailable");
+		this.connection = connection;
+	}
+	
+	public DstoreConnection getConnection() {
+		return connection;
+	}
+}

Dstore.java

@@ -335,7 +335,7 @@ public class Dstore {
 							messageReceived(socket, receivedMessage);
 							if(!receivedMessage.equals(Protocol.ACK_TOKEN)) {
 								//Log error
-								System.out.println("Dstore " + dstore + " should have sent ACK but " + port + " received " + receivedMessage);
+								System.err.println("Dstore " + dstore + " should have sent ACK but " + port + " received " + receivedMessage);
 							}
 							
 							byte[] content = new byte[BUFFER_SIZE];
@@ -378,7 +378,6 @@ public class Dstore {
 				controllerOut.println(Protocol.REBALANCE_COMPLETE_TOKEN);
 				messageSent(controllerSocket, Protocol.REBALANCE_COMPLETE_TOKEN);
 			}
-			System.out.println("Sent message REBALANCE_COMPLETE");
 		}).start();
 	}
 	

DstoreConnection.java

@@ -52,12 +52,16 @@ public class DstoreConnection {
 		return new DstoreDisconnectException(this);
 	}
 	
-	public String sendAndReceive(String message, String... expectedMessages) throws DstoreDisconnectException {
+	public void checkAvailable() throws DeadStoreException {
+		if(!available) throw new DeadStoreException(this);
+	}
+	
+	public String sendAndReceive(String message, String... expectedMessages) throws DstoreDisconnectException, DeadStoreException {
 		System.out.println("Getting lock...");
 		synchronized(this) {
 			try {
 				System.out.println("Lock acquired");
-				if(!available) throw getDisconnectData();
+				checkAvailable();
 				writer.println(message);
 				writer.flush();
 				//System.out.println("Controller sent " + message + " to port " + port);
@@ -72,7 +76,7 @@ public class DstoreConnection {
 		}
 	}
 	
-	public String receive(String... expectedMessages) throws DstoreDisconnectException {
+	public String receive(String... expectedMessages) throws DstoreDisconnectException, DeadStoreException {
 		String findMessage = checkQueue(expectedMessages);
 		if(findMessage != null) {
 			return findMessage;
@@ -81,7 +85,7 @@ public class DstoreConnection {
 		System.out.println("Getting lock...");
 		synchronized(this) {
 			System.out.println("Lock acquired");
-			if(!available) throw getDisconnectData();
+			checkAvailable();
 			
 			//Check the queue twice: once incase the receiver is busy, twice incase the message was added by the last thread
 			findMessage = checkQueue(expectedMessages);
@@ -130,36 +134,12 @@ public class DstoreConnection {
 			e.printStackTrace();
 		}
 		
-			/*
-			String returnMessage = null;
-			do {
-				returnMessage = reader.readLine();
-				if(returnMessage == null) {
-					System.out.println("Dstore disconnected");
-					available = false;
-					throw new DstoreDisconnectException();
-				}
-				if(expectedMessage != null && !expectedMessage.equals(returnMessage.split(" ")[0])) {
-					queue.add(returnMessage);
-					if(queue.size() > MAX_QUEUE_SIZE) queue.remove(0);
-					returnMessage = null;
-				}
-			}
-			while(returnMessage == null);
-			System.out.println("Controller received " + returnMessage);
-			return returnMessage;
-		}
-		catch(IOException e) {
-			e.printStackTrace();
-			return "";
-		}
-		*/
 		return "";
 	}
 	
 	protected boolean isExpected(String message, String[] expectedMessages) {
 		for(String s : expectedMessages) {
-			if(s.equals(message.split(" ")[0])) return true;
+			if(s.equals(message)) return true;
 		}
 		return false;
 	}