waf/playground/netcache/Netcache.java

// Thomas Nagy, 2011

// TODO handle all exceptions properly

import java.util.HashMap;
import java.util.Map;
import java.util.List;
import java.util.Date;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.Collections;
import java.lang.Math;
import java.lang.StringBuilder;
import java.io.*;
import java.net.*;
import java.security.*;

public class Netcache implements Runnable, Comparator<Object[]> {
	private static int PORT = 51200;
	private static String CACHEDIR = "/tmp/wafcache/";
	private static long MAX = 10l * 1024l * 1024l * 1024l;
	private static double CLEANRATIO = 0.8;
	private static int BUF = 16 * 8192;

	private static HashMap<String, Object[]> flist = null;
	private Socket server = null;

	public Netcache(Socket server) {
		this.server = server;
	}

	public void run () {
		try {

			while (true) {
				InputStream in = server.getInputStream();
				OutputStream out = server.getOutputStream();

				byte b[] = new byte[128];
				int off = 0;
				while (off < b.length) {
					off += in.read(b, off, b.length - off);
				}

				//System.out.println(new String(b));
				String[] args = new String(b).split(",");
				if (args[0].equals("LST")) {
					lst(args, in, out);
				}
				else if (args[0].equals("PUT")) {
					put(args, in, out);
				}
				else if (args[0].equals("GET")) {
					get(args, in, out);
				}
				else if (args[0].equals("CLEAN")) {
					clean(args, in, out);
				}
				else if (args[0].equals("BYE")) {
					server.close();
					break;
				}
				else {
					System.out.println("Invalid command " + new String(b));
					server.close();
					break;
				}
			}
		} catch (IOException e) {
			e.printStackTrace();
		}
	}

	public void lst(String[] args, InputStream in, OutputStream out) throws IOException {
		StringBuilder b = new StringBuilder();
		int k = 0;
		synchronized(flist) {
			for (String name : flist.keySet()) {
				b.append(name);
				if (k <= flist.size()) {
					k++;
					b.append("\n");
				}
			}
		}

		byte[] ret = b.toString().getBytes();
		String header = String.format("%-128s", String.format("%d,", ret.length));

		out.write(header.getBytes());
		out.write(ret);
	}

	public void put(String[] args, InputStream in, OutputStream out) throws IOException {
		File cachedir = new File(CACHEDIR);
		File temp = File.createTempFile("foo", ".suffix", cachedir);

		long size = new Long(args[3].trim());

		//System.out.println("" + args[1] + " " + args[2] + " " + args[3] + " " + args.length);

		byte[] buf = new byte[BUF];
		long cnt = 0;
		OutputStream w = new FileOutputStream(temp);
		try {
			while (cnt < size) {
				int c = in.read(buf, 0, (int) Math.min(BUF, size-cnt));
				if (c == 0) {
					throw new RuntimeException("Connection closed too early");
				}
				w.write(buf, 0, c);
				cnt += c;
			}
		} finally {
			w.close();
		}

		/*if (cnt != size) {
		  System.out.println("error!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!");
		  }*/

		File parent = new File(new File(new File(CACHEDIR), args[1].substring(0, 2)), args[1]);
		File dest = new File(parent, args[2]);
		try {
			dest.getParentFile().mkdirs();
		} catch (Exception e) {
		}

		if (!temp.renameTo(dest)) {
			throw new RuntimeException("Could not reanae the filee");
		}

		long total = 0;
		for (File f : parent.listFiles()) {
			total += f.length();
		}

		synchronized(flist) {
			if (flist.containsKey(parent.getName())) {
				flist.get(parent.getName())[0] = parent.lastModified();
			}
			else
			{
				flist.put(parent.getName(), new Object[] {parent.lastModified(), total, parent.getName()});
			}
		}
	}

	public void get(String[] args, InputStream in, OutputStream out) throws IOException {
		File f = new File(new File(new File(new File(CACHEDIR), args[1].substring(0, 2)), args[1]), args[2].trim());
		long fsize = -1;
		try {
			fsize = f.length();
		} catch (Exception e) {
			// return -1 to the client
		}

		String ret = String.format("%-128s", String.format("%d,", fsize));
		out.write(ret.getBytes());

		byte[] buf = new byte[BUF];

		long cnt = 0;
		InputStream s = new FileInputStream(f);
		try {
			while (cnt < fsize) {
				long c = s.read(buf);
				cnt += c;
				out.write(buf, 0, (int) c);
			}
		} finally {
			s.close();
		}

		File parent = f.getParentFile();
		Date d = new Date();
		parent.setLastModified(d.getTime());
		synchronized(flist) {
			flist.get(parent.getName())[0] = parent.lastModified();
		}
	}

	public void clean(String[] args, InputStream in, OutputStream out) throws IOException {
		synchronized(flist) {
			long total = 0;
			for (Map.Entry<String, Object[]> entry : flist.entrySet()) {
				total += (Long) entry.getValue()[1];
			}

			List<Object[]> k = new ArrayList<Object[]>(flist.values());
			Collections.sort(k, this);

			int cur = 0;
			while (total > MAX * CLEANRATIO) {
				Object[] kk = k.get(cur);

				String name = (String) kk[2];
				File f = new File(new File(new File(CACHEDIR), name.substring(0, 2)), name);
				//System.out.println("removing " + cur + " " + kk[0] + " " + kk[1] + " " + f.getAbsolutePath());
				rm(f);

				total -= (Long) kk[1];

				flist.remove(name);
				cur++;
			}
		}
	}

	public static void init_flist() {
		flist = new HashMap<String, Object[]>();
		synchronized(flist) {
			File dir = new File(CACHEDIR);
			try {
				dir.mkdirs();
			} catch (Exception e) {

			}

			for (File d : dir.listFiles()) {
				if (!d.isDirectory()) continue;
				for (File sd : d.listFiles()) {
					if (!sd.isDirectory()) continue;
					long total = 0;
					for (File f : sd.listFiles()) {
						total += f.length();
					}
					//System.out.println(sd.getName());
					flist.put(sd.getName(), new Object[] {sd.lastModified(), total, sd.getName()});
				}
			}
		}
	}

	public int compare(Object[] a, Object[] b) {
		return ((Long) a[0]).compareTo((Long) b[0]);
	}

	public static void rm(File dir) {
		if (dir.isDirectory()) {
			for (File f: dir.listFiles())
			{
				rm(f);
			}
		}
		dir.delete();
	}

	public static void main(String[] args) {
		try
		{
			init_flist();
			System.out.println("ready (" + flist.keySet().size() + " dirs)");
			ServerSocket sock = new ServerSocket(PORT);
			while (true)
			{
				Netcache tmp = new Netcache(sock.accept());
				Thread t = new Thread(tmp);
				t.start();
			}
		} catch (IOException e) {
			e.printStackTrace();
		}
	}
}