java DiskQueue

package com.sibo.concurrent;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.*;
import java.nio.file.*;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;


public class DiskQueue {

    private final static Logger LOG = LoggerFactory.getLogger(DiskQueue.class);
    private final static int MAX_QUEUE_SIZE = 1000;
    private final static long DEFAULT_SYNC_EVERY = 2 * 60 * 1000L; // 2 minute

    // run-time state (also persisted to disk)
    private long readPos;
    private long writePos;
    private long readFileNum;
    private long writeFileNum;
    private AtomicLong depth = new AtomicLong(0L);

    private Lock lock = new ReentrantLock();

    // instantiation time metadata
    private final String name;
    private final String dataPath;
    private final long maxBytesPerFile;
    private final long syncEvery;
    private volatile boolean needSync;
    private volatile boolean exit;

    private long lastWriteTime;

    // keeps track of the positon where we have read
    private long nextReadPos;
    private long nextReadFileNum;

    private RandomAccessFile readFile;
    private RandomAccessFile writeFile;

    //private BlockingQueue<byte[]> readQueue;

    private ArrayBlockingQueue<byte[]> writeQueue;

    public DiskQueue(String name, String dataPath, int maxBytesPerFile) {
        this(name, dataPath, maxBytesPerFile, DEFAULT_SYNC_EVERY);
    }

    public DiskQueue(String name, String dataPath, int maxBytesPerFile, long syncEvery) {
        this.name = name;
        this.dataPath = dataPath;
        this.maxBytesPerFile = maxBytesPerFile;
        this.syncEvery = syncEvery;
        this.needSync = true;
        this.exit = false;
        this.writeQueue = new ArrayBlockingQueue<byte[]>(MAX_QUEUE_SIZE);
        this.lastWriteTime = System.currentTimeMillis();
        //this.readQueue = new ArrayBlockingQueue<byte[]>(DEFAULT_MAX_QUEUE_SIZE);
        this.retrieveMetaData();
    }

    public void put(byte[] bytes) {
        if (this.exit) return;

        long curTime = System.currentTimeMillis();
        if ((curTime - this.lastWriteTime >= this.syncEvery && !this.writeQueue.isEmpty()) || this.writeQueue.size() >= MAX_QUEUE_SIZE) {
            this.lastWriteTime = curTime;
            writeToFile();
        }
        if (this.needSync) {
            this.sync();
        }
        try {
            this.writeQueue.put(bytes);
        } catch (InterruptedException e) {
            LOG.error("diskqueue({}), put element to writeQueue failed, {}", this.name, e.getMessage());
        }
    }

    public void close() {
        try {
            lock.lock();
            this.exit = true;
            this.writeToFile();
            if (this.readFile != null) {
                this.readFile.close();
                this.readFile = null; // help gc
            }
            if (this.writeFile != null) {
                this.writeFile.close();
                this.writeFile = null; // help gc
            }
            this.sync();
            //this.executorService.shutdown();
        } catch (IOException e) {
            LOG.error("diskqueue({}) close failed : {}", this.name, e.getMessage());
        } finally {
            lock.unlock();
        }
    }

    private void deleteAllFiles() {
        this.skipToNextRWFile();
        Path file = Paths.get(this.metaDataFileName());
        try {
            if(Files.exists(file, LinkOption.NOFOLLOW_LINKS) && !Files.deleteIfExists(file)) {
                LOG.error("diskqueue({}), delete metaFile failed, {}", this.name, file.getFileName().toString());
            }
        } catch (IOException e) {
            e.printStackTrace();
            LOG.error("diskqueue({}), delete metaFile {}, exception {}", this.name, file.getFileName().toString(), e.getMessage());
        }
    }

    private void skipToNextRWFile() {
        if (this.readFile != null) {
            try {
                this.readFile.close();
            } catch (IOException e) {
                LOG.error("diskqueue({}), failed to close readFile, {}", this.name, e.getMessage());
            }
            this.readFile = null;
        }
        if (this.writeFile != null) {
            try {
                this.writeFile.close();
            } catch (IOException e) {
                LOG.error("diskqueue({}), failed to close writeFile, {}", this.name, e.getMessage());
            }
            this.writeFile = null;
        }
        for (long i = this.readFileNum; i <= this.writeFileNum; i++) {
            Path file = Paths.get(this.fileName(i));
            try {
                if (Files.exists(file, LinkOption.NOFOLLOW_LINKS)) {
                    Files.delete(file);
                } else {
                    LOG.error("diskqueue({}), delete file failed, {}", this.name, file.getFileName().toString());
                }
            } catch (IOException e) {
                e.printStackTrace();
                LOG.error("diskqueue({}), delete file {}, exception {}", this.name, file.getFileName().toString(), e.getMessage());
            }
        }
        this.writeFileNum++;
        this.writePos = 0;
        this.readFileNum = this.writeFileNum;
        this.readPos = 0;
        this.nextReadFileNum = this.writeFileNum;
        this.nextReadPos = 0;
        this.depth.set(0L);
    }

    public byte[] readOne() {
        if (this.readFile == null) {
            String curFileName = this.fileName(this.readFileNum);
            try {
                Path file = Paths.get(curFileName);
                if (!Files.exists(file, LinkOption.NOFOLLOW_LINKS)) {
                    Files.createFile(file);
                }
                this.readFile = new RandomAccessFile(file.toFile(), "r");
                LOG.info("disqueue({}): readOne() opened {}", this.name, curFileName);
                if (this.readPos > 0) {
                    this.readFile.seek(this.readPos);
                }
            } catch (IOException e) {
                if (this.readFile != null) {
                    try {
                        this.readFile.close();
                    } catch (IOException e1) {
                        e1.printStackTrace();
                    }
                }
                this.readFile = null;
                LOG.error("open file: {} file fail", curFileName);
                return null;
            }
        }

        try {
            int dataLen = this.readFile.readInt();
            byte[] result = new byte[dataLen];
            this.readFile.read(result);
            this.nextReadPos = this.readPos + (4 + dataLen);
            this.nextReadFileNum = this.readFileNum;
            if (this.nextReadPos > this.maxBytesPerFile) {
                if (this.readFile != null) {
                    this.readFile.close();
                    this.readFile = null;
                }
                this.nextReadFileNum++;
                this.readPos = 0;
            }
            return result;
        } catch (IOException e) {
            if (this.readFile != null) {
                try {
                    this.readFile.close();
                } catch (IOException e1) {
                    e1.printStackTrace();
                }
            }
            this.readFile = null;
        }
        return null;
    }

    private void writeOne(byte[] data) {
        if (this.writeFile == null) {
            String curFileName = this.fileName(this.writeFileNum);
            try {
                Path file = Paths.get(curFileName);
                //File file = new File(curFileName);
                if (!Files.exists(file, LinkOption.NOFOLLOW_LINKS)) {
                    Files.createFile(file);
                }
                this.writeFile = new RandomAccessFile(file.toFile(), "rw");
                LOG.info("disqueue({}): writeOne() opened {}", this.name, curFileName);
                if (this.writePos > 0) {
                    this.writeFile.seek(this.writePos);
                }
            } catch (IOException e) {
                if (this.writeFile != null) {
                    try {
                        this.writeFile.close();
                    } catch (IOException e1) {
                        e1.printStackTrace();
                    }
                }
                this.writeFile = null;
                LOG.error("open file: {} file fail", curFileName);
                return;
            }
        }
        try {
            int dataLen = data.length;
            this.writeFile.writeInt(dataLen);
            this.writeFile.write(data);
            this.writePos += (4 + data.length);
            this.depth.incrementAndGet();

            if (this.writePos > this.maxBytesPerFile) {
                this.writeFileNum++;
                this.writePos = 0;
                this.sync();
                if (this.writeFile != null) {
                    this.writeFile.close();
                    this.writeFile = null;
                }
            }
        } catch (IOException e) {
            if (this.writeFile != null) {
                try {
                    this.writeFile.close();
                } catch (IOException e1) {
                    e1.printStackTrace();
                }
            }
            this.writeFile = null;
        }
    }

    private void sync() {
        try {
            lock.lock();
            if (this.writeFile != null) {
                FileDescriptor fd = this.writeFile.getFD();
                fd.sync();
            }
            this.persistMetaData();
            this.needSync = false;
        } catch (IOException e) {
            //e.printStackTrace();
            LOG.error("diskqueue({}), failed to sync, {}", this.name, e.getMessage());
        } finally {
            lock.unlock();
        }
    }

    private void checkTailCorruption(long depth) {
        if (this.readFileNum < this.writeFileNum || this.readPos < this.writePos) {
            return;
        }
        if (this.depth.get() != 0) {
            if (depth < 0) {
                LOG.error("diskqueue {} negative depth at tail {}, meta data corruption, resetting 0...", this.name, depth);
            } else if (depth > 0) {
                LOG.error("diskqueue {} positive depth at tail {}, data loss, resetting 0...", this.name, depth);
            }
            this.depth.set(0L);
            this.needSync = true;
        }

        if (this.readFileNum != this.writeFileNum || this.readPos != this.writePos) {
            if (this.readFileNum > this.writeFileNum) {
                LOG.error("disqueue({}) readFileNum > writeFileNum ({} > {}), corruption, skipping to next writeFileNum and resetting 0...", this.name, this.readFileNum, this.writeFileNum);
            }
            if (this.readPos > this.writePos) {
                LOG.error("diskqueue({}) readPos > writePos ({} > {}), corruption, skipping to next writeFileNum and resetting 0...", this.name, this.readPos, this.writePos);
            }
            this.skipToNextRWFile();
            this.needSync = true;
        }
    }

    // move forward to next read file
    private void moveForward() {
        long oldReadFileNum = this.readFileNum;
        this.readFileNum = this.nextReadFileNum;
        this.readPos = this.nextReadPos;
        long depth = this.depth.decrementAndGet();

        if (oldReadFileNum != this.nextReadFileNum) {
            this.needSync = true;
            String fileName = this.fileName(oldReadFileNum);
            try {
                Path file = Paths.get(fileName);
                if (!Files.deleteIfExists(file)) {
                    LOG.error("diskqueue({}) failed to Remove({})", this.name, fileName);
                }
            } catch (IOException e) {
                e.printStackTrace();
                LOG.error("diskqueue({}) failed to Remove({}), exception {}", this.name, fileName, e.getMessage());
            }
        }
        this.checkTailCorruption(depth);
    }

    private void handleReadError() {
        try {
            lock.lock();
            if (this.readFileNum == this.writeFileNum) {
                if (this.writeFile != null) {
                    this.writeFile.close();
                    this.writeFile = null;
                }
                this.writeFileNum++;
                this.writePos = 0;
            }
            String fileName = this.fileName(this.readFileNum);
            String badFileName = fileName + ".bad";
            LOG.warn("diskqueue({}) jump to next file and saving bad file as {}", this.name, badFileName);
            Path file = Paths.get(fileName);
            Path targetFile = Paths.get(badFileName);
            Files.move(file, targetFile, StandardCopyOption.REPLACE_EXISTING);
            this.readFileNum++;
            this.readPos = 0;
            this.nextReadFileNum = this.readFileNum;
            this.nextReadPos = 0;
            this.needSync = true;

        } catch (IOException e) {
            LOG.error("diskqueue({}) failed to rename bad diskqueue file {} to {}", this.name, this.fileName(this.readFileNum), this.fileName(this.readFileNum) + ".bad");
        } finally {
            lock.unlock();
        }
    }

    private String metaDataFileName() {
        return this.dataPath + File.separator + this.name + ".diskqueue.meta.dat";
    }

    private String fileName(long fileNum) {
        return this.dataPath + File.separator + this.name + ".diskqueue." + fileNum + ".dat";
    }

    // init metadata from meta file
    private void retrieveMetaData() {
        String fileName = this.metaDataFileName();
        try {
            //File file = new File(fileName);
            Path file = Paths.get(fileName);
            if (Files.exists(file, LinkOption.NOFOLLOW_LINKS)) {
                DataInputStream dis = new DataInputStream(new FileInputStream(file.toFile()));
                this.depth.set(dis.readLong());
                this.readFileNum = dis.readLong();
                this.readPos = dis.readLong();
                this.writeFileNum = dis.readLong();
                this.writePos = dis.readLong();
                this.nextReadFileNum = this.readFileNum;
                this.nextReadPos = this.readPos;
                dis.close();
            }
        } catch (IOException e) {
            LOG.error("diskqueue({}) open metaData file {} failed, ", this.name, fileName, e.getMessage());
        }
    }

    // persist metadata to disk meta file
    private void persistMetaData() {
        //Random random = new
        String fileName = this.metaDataFileName();
        try {
            Path source = Paths.get(fileName + "_" + System.nanoTime() + ".tmp");
            //Path source = Files.createTempFile(this.simpleMetaDataFileName() + ((int) (Math.random() * Integer.MAX_VALUE)), ".tmp");
            if (!Files.exists(source, LinkOption.NOFOLLOW_LINKS)) {
                Files.createFile(source);
            }
            //File file = source.toFile();
            DataOutputStream dos = new DataOutputStream(new FileOutputStream(source.toFile()));
            //FileDescriptor fd = fis.getFD();
            dos.writeLong(this.depth.get());
            dos.writeLong(this.readFileNum);
            dos.writeLong(this.readPos);
            dos.writeLong(this.writeFileNum);
            dos.writeLong(this.writePos);
            dos.flush();
            //fd.sync();
            dos.close();
            Files.move(source, Paths.get(fileName), StandardCopyOption.REPLACE_EXISTING);
        } catch (IOException e) {
            LOG.error("diskqueue({}), failed to persist metaDataFile {}: {}", this.name, fileName, e.getMessage());
        }
    }

    // todo 待优化，需考虑文件过多，数据量过大的情况(通常是断线太久导致数据堆积)
    // optimize, consider when the file is too big or too many files, this offen happens when the remote server close too long
    public List<byte[]> readAllBytesFromFile() {
        List<byte[]> bytesList = new ArrayList<>();

        // get data from the current writeQueue
        try {
            while (!this.writeQueue.isEmpty()) {
                bytesList.add(this.writeQueue.take());
            }
        } catch (InterruptedException e) {
            LOG.error("diskqueue({}), take element form writeQueue failed, {}", this.name, e.getMessage());
        }

        // get data from disk files
        while (this.readFileNum < this.writeFileNum || this.readPos < this.writePos) {
            if (this.nextReadPos == this.readPos) {
                byte[] bytes = this.readOne();
                if (bytes == null) {
                    this.handleReadError();
                } else {
                    bytesList.add(bytes);
                    this.moveForward();
                }
            }
        }
        this.sync();
        return bytesList;
    }

    // write data to disk file
    private void writeToFile() {
        try {
            lock.lock();
            while (!this.writeQueue.isEmpty()) {
                this.writeOne(this.writeQueue.take());
            }
            this.sync();
        } catch (InterruptedException e) {
            LOG.error("diskqueue({}) writeQueue take failed : {}", this.name, e.getMessage());
        } finally {
            lock.unlock();
        }
    }


}