package com.github.zhongl.jtoolkit;

import java.io.*;
import java.nio.ByteBuffer;
import java.nio.channels.FileChannel;

/**
 * IDPoolBaseOnFile is a pool for reusing a set of num id, and can be persisted in a file.
 *
 * @author <a href="mailto:zhong.lunfu@gmail.com">zhongl</a>
 */
public class IDPoolBaseOnFile {
  private final int capacity;
  private final RandomAccessFile raf;
  private final ByteBuffer bits;

  private int inUsed = 0;

  public IDPoolBaseOnFile(final int capacity, final String file) throws IOException {
    this.capacity = capacity;
    raf = new RandomAccessFile(file, "rwd");
    try {
      setFileLength(alignAt8(capacity));
      bits = raf.getChannel().map(FileChannel.MapMode.READ_WRITE, 0, raf.length());
    } catch (IOException e) {
      dispose();
      throw e;
    }
  }

  public static void main(String... args) throws Exception {
    final String file = args[0];
    final int size = Integer.parseInt(args[1]);
    final IDPoolBaseOnFile pool = new IDPoolBaseOnFile(size, file);

    long elapse = 0;

    for (int i = 0; i < size; i++) {
      long begin = System.nanoTime();
      pool.acquire();
      elapse += System.nanoTime() - begin;
    }

    System.out.println(elapse / size);
    elapse = 0;

    for (int i = 0; i < pool.capacity; i++) {
      long begin = System.nanoTime();
      pool.release(i);
      elapse += System.nanoTime() - begin;
    }
    System.out.println(elapse / size);

    System.out.println(pool.bits.equals(ByteBuffer.wrap(new byte[alignAt8(size)])));
    pool.dispose();

  }

  public synchronized void dispose() { if (raf != null) try { raf.close(); } catch (IOException e) {} }

  /**
   * Release id to pool.
   *
   * @param id
   */
  public synchronized void release(int id) { if (set(id, false)) inUsed--; }

  /**
   * Acquire a id from pool.
   *
   * @return id
   */
  public synchronized int acquire() {
    if (inUsed == capacity) throw new IllegalStateException("No more id for acquisition.");
    for (; ;) {
      final int id = nextId();
      final byte cur = bits.get(index(id));
      final byte or = (byte) (cur | bit(id));
      if (or == cur || id >= capacity) continue; // acquired id
      bits.put(index(id), or);
      inUsed++;
      return id;
    }
  }

  /**
   * Acquire specify id.
   *
   * @param id
   */
  public synchronized void acquire(int id) { if (set(id, true)) inUsed++; }

  /** Reset state to init. */
  public synchronized void reset() {
    inUsed = 0;
    for (int i = 0; i < capacity; i++) release(i);
  }

  private static int alignAt8(final int num) {return num / 8 + (num % 8 == 0 ? 0 : 1);}

  private static int bit(long id) { return (1 << id % 8); }

  private static int index(int id) { return id / 8; }

  private void setFileLength(int length) throws IOException {if (raf.length() < length) raf.setLength(length);}

  /**
   * Set index bit to 0 or 1.
   *
   * @param id
   * @param b  true for aquire(1), false for release(0).
   *
   * @return false means no change.
   */
  private boolean set(int id, boolean b) {
    if (id < 0 || id >= capacity) throw new IllegalArgumentException("id : " + id);
    final byte o = bits.get(index(id));
    final byte n = b ? (byte) (o | bit(id)) : (byte) (o & ~bit(id));
    bits.put(index(id), n);
    return o != n;
  }

  private int nextId() { return (int) (System.nanoTime() % capacity); }
}

package com.github.zhongl.jtoolkit;

import java.util.BitSet;

import org.junit.*;

public class IDPoolBaseOnFileTest {
  private static final int CAPACITY = 1000;
  private static IDPoolBaseOnFile pool;

  @BeforeClass
  public static void setUpClass() throws Exception { pool = new IDPoolBaseOnFile(CAPACITY, "target/id.set"); }

  @Test(expected = IllegalArgumentException.class)
  public void illegalIdLessThanZero() throws Exception { pool.acquire(-1); }

  @Test(expected = IllegalArgumentException.class)
  public void illegalIdEqualsCapacity() throws Exception { pool.acquire(CAPACITY); }

  @Test(expected = IllegalStateException.class)
  public void noMoreId() throws Exception {
    for (int i = 0; i < CAPACITY; i++) pool.acquire(i);
    pool.acquire();
  }

  @Test
  public void acquiredValidId() throws Exception {
    final BitSet ids = new BitSet(CAPACITY);
    for (int i = 0; i < CAPACITY; i++) {
      final int id = pool.acquire();
      if (id > 999 || id < 0 || ids.get(id)) throw new IllegalStateException("invalid id :" + id);
      ids.set(id);
    }
  }

  @After
  public void tearDown() throws Exception { pool.reset(); }

  @AfterClass
  public static void tearDownClass() throws Exception { if (pool != null) pool.dispose(); }
}