<commit_before>package org.yamcs.yarch.rocksdb;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Iterator;
import java.util.NoSuchElementException;
import java.util.PriorityQueue;
import org.rocksdb.RocksDBException;
import org.slf4j.Logger;
import org.yamcs.utils.ByteArrayUtils;
import org.yamcs.utils.LoggingUtils;
import org.yamcs.utils.TimeInterval;
import org.yamcs.yarch.HistogramRecord;
import org.yamcs.yarch.HistogramSegment;
import org.yamcs.yarch.PartitionManager;
import org.yamcs.yarch.TableDefinition;
import org.yamcs.yarch.YarchDatabaseInstance;
import static org.yamcs.yarch.HistogramSegment.segmentStart;

import static org.yamcs.yarch.rocksdb.RdbStorageEngine.TBS_INDEX_SIZE;

/**
 * 
 * @author nm
 *
 */
class RdbHistogramIterator implements Iterator<HistogramRecord> {

    private Iterator<PartitionManager.Interval> intervalIterator;
    private AscendingRangeIterator segmentIterator;

    private PriorityQueue<HistogramRecord> records = new PriorityQueue<>();

    private final TimeInterval interval;
    private final long mergeTime;
    private final Tablespace tablespace;

    YarchDatabaseInstance ydb;
    TableDefinition tblDef;
    YRDB rdb;

    Logger log;
    String colName;
    boolean stopReached = false;

    // FIXME: mergeTime does not merge records across partitions or segments
    public RdbHistogramIterator(Tablespace tablespace, YarchDatabaseInstance ydb, TableDefinition tblDef,
            String colName, TimeInterval interval, long mergeTime) throws RocksDBException, IOException {
        this.interval = interval;
        this.mergeTime = mergeTime;
        this.ydb = ydb;
        this.tblDef = tblDef;
        this.colName = colName;
        this.tablespace = tablespace;

        PartitionManager partMgr = RdbStorageEngine.getInstance().getPartitionManager(tblDef);
        intervalIterator = partMgr.intervalIterator(interval);
        log = LoggingUtils.getLogger(this.getClass(), ydb.getName(), tblDef);
        readNextPartition();
    }

    private void readNextPartition() throws RocksDBException, IOException {
        if(!intervalIterator.hasNext()) {
            stopReached = true;
            return;
        }

        PartitionManager.Interval intv = intervalIterator.next();
        
        if (rdb != null) {
            tablespace.dispose(rdb);
        }

        RdbHistogramInfo hist = (RdbHistogramInfo) intv.getHistogram(colName);

        if(hist==null) {
            readNextPartition();
            return;
        }
        rdb = tablespace.getRdb(hist.partitionDir, false);
        
        long segStart = interval.hasStart() ? segmentStart(interval.getStart()) : 0;
        byte[] dbKeyStart = ByteArrayUtils.encodeInt(hist.tbsIndex, new byte[12], 0);
        ByteArrayUtils.encodeLong(segStart, dbKeyStart, TBS_INDEX_SIZE);

        boolean strictEnd;
        byte[] dbKeyStop;
        if (interval.hasEnd()) {
            strictEnd = false;
            dbKeyStop = ByteArrayUtils.encodeInt(hist.tbsIndex, new byte[12], 0);
            long segStop = segmentStart(interval.getEnd());
            ByteArrayUtils.encodeLong(segStop, dbKeyStop, TBS_INDEX_SIZE);
        } else {
            dbKeyStop = ByteArrayUtils.encodeInt(hist.tbsIndex + 1, new byte[12], 0);
            strictEnd = true;
        }
        
        segmentIterator = new AscendingRangeIterator(rdb.newIterator(), dbKeyStart, false, dbKeyStop, strictEnd);
        readNextSegments();
        tablespace.dispose(rdb);
    }

    // reads all the segments with the same sstart time
    private void readNextSegments() throws RocksDBException, IOException {
        if (!segmentIterator.isValid()) {
            readNextPartition();
            return;
        }

        ByteBuffer bb = ByteBuffer.wrap(segmentIterator.key());
        long sstart = bb.getLong(RdbStorageEngine.TBS_INDEX_SIZE);
        
        while (true) {
            boolean beyondStop = addRecords(segmentIterator.key(), segmentIterator.value());
            if (beyondStop) {
                stopReached = true;
            }

            segmentIterator.next();
            if (!segmentIterator.isValid()) {
                readNextPartition();
                break;
            }
            bb = ByteBuffer.wrap(segmentIterator.key());
            long g = bb.getLong();
            if (g != sstart) {
                break;
            }
        }
    }

    public void close() {
        if (rdb != null) {
            tablespace.dispose(rdb);
            rdb = null;
        }
    }

    // add all records from this segment into the queue
    // if the stop has been reached add only partially the records, return true
    private boolean addRecords(byte[] key, byte[] val) {
        ByteBuffer kbb = ByteBuffer.wrap(key, TBS_INDEX_SIZE, key.length-TBS_INDEX_SIZE);
        long sstart = kbb.getLong();
        byte[] columnv = new byte[kbb.remaining()];
        kbb.get(columnv);
        ByteBuffer vbb = ByteBuffer.wrap(val);
        HistogramRecord r = null;
        while (vbb.hasRemaining()) {
            long start = sstart * HistogramSegment.GROUPING_FACTOR + vbb.getInt();
            
            long stop = sstart * HistogramSegment.GROUPING_FACTOR + vbb.getInt();
            int num = vbb.getShort();
            if ((interval.hasStart()) && (stop < interval.getStart())) {
                continue;
            }
            if ((interval.hasEnd()) && (start > interval.getEnd())) {
                if (r != null) {
                    records.add(r);
                }
                return true;
            }
            if (r == null) {
                r = new HistogramRecord(columnv, start, stop, num);
            } else {
                if (start - r.getStop() < mergeTime) {
                    r = new HistogramRecord(r.getColumnv(), r.getStart(), stop, r.getNumTuples() + num);
                } else {
                    records.add(r);
                    r = new HistogramRecord(columnv, start, stop, num);
                }
            }
        }
        if (r != null) {
            records.add(r);
        }
        return false;<commit_msg>bugfix: close the iterator to avoid resource leaks<commit_after>package org.yamcs.yarch.rocksdb;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Iterator;
import java.util.NoSuchElementException;
import java.util.PriorityQueue;
import org.rocksdb.RocksDBException;
import org.slf4j.Logger;
import org.yamcs.utils.ByteArrayUtils;
import org.yamcs.utils.LoggingUtils;
import org.yamcs.utils.TimeInterval;
import org.yamcs.yarch.HistogramRecord;
import org.yamcs.yarch.HistogramSegment;
import org.yamcs.yarch.PartitionManager;
import org.yamcs.yarch.TableDefinition;
import org.yamcs.yarch.YarchDatabaseInstance;
import static org.yamcs.yarch.HistogramSegment.segmentStart;

import static org.yamcs.yarch.rocksdb.RdbStorageEngine.TBS_INDEX_SIZE;

/**
 * 
 * @author nm
 *
 */
class RdbHistogramIterator implements Iterator<HistogramRecord> {

    private Iterator<PartitionManager.Interval> intervalIterator;
    private AscendingRangeIterator segmentIterator;

    private PriorityQueue<HistogramRecord> records = new PriorityQueue<>();

    private final TimeInterval interval;
    private final long mergeTime;
    private final Tablespace tablespace;

    YarchDatabaseInstance ydb;
    TableDefinition tblDef;
    YRDB rdb;

    Logger log;
    String colName;
    boolean stopReached = false;

    // FIXME: mergeTime does not merge records across partitions or segments
    public RdbHistogramIterator(Tablespace tablespace, YarchDatabaseInstance ydb, TableDefinition tblDef,
            String colName, TimeInterval interval, long mergeTime) throws RocksDBException, IOException {
        this.interval = interval;
        this.mergeTime = mergeTime;
        this.ydb = ydb;
        this.tblDef = tblDef;
        this.colName = colName;
        this.tablespace = tablespace;

        PartitionManager partMgr = RdbStorageEngine.getInstance().getPartitionManager(tblDef);
        intervalIterator = partMgr.intervalIterator(interval);
        log = LoggingUtils.getLogger(this.getClass(), ydb.getName(), tblDef);
        readNextPartition();
    }

    private void readNextPartition() throws RocksDBException, IOException {
        if (!intervalIterator.hasNext()) {
            stopReached = true;
            return;
        }

        PartitionManager.Interval intv = intervalIterator.next();

        if (rdb != null) {
            tablespace.dispose(rdb);
        }

        RdbHistogramInfo hist = (RdbHistogramInfo) intv.getHistogram(colName);

        if (hist == null) {
            readNextPartition();
            return;
        }
        rdb = tablespace.getRdb(hist.partitionDir, false);

        long segStart = interval.hasStart() ? segmentStart(interval.getStart()) : 0;
        byte[] dbKeyStart = ByteArrayUtils.encodeInt(hist.tbsIndex, new byte[12], 0);
        ByteArrayUtils.encodeLong(segStart, dbKeyStart, TBS_INDEX_SIZE);

        boolean strictEnd;
        byte[] dbKeyStop;
        if (interval.hasEnd()) {
            strictEnd = false;
            dbKeyStop = ByteArrayUtils.encodeInt(hist.tbsIndex, new byte[12], 0);
            long segStop = segmentStart(interval.getEnd());
            ByteArrayUtils.encodeLong(segStop, dbKeyStop, TBS_INDEX_SIZE);
        } else {
            dbKeyStop = ByteArrayUtils.encodeInt(hist.tbsIndex + 1, new byte[12], 0);
            strictEnd = true;
        }
        if (segmentIterator != null) {
            segmentIterator.close();
        }

        segmentIterator = new AscendingRangeIterator(rdb.newIterator(), dbKeyStart, false, dbKeyStop, strictEnd);
        readNextSegments();

    }

    // reads all the segments with the same sstart time
    private void readNextSegments() throws RocksDBException, IOException {
        if (!segmentIterator.isValid()) {
            readNextPartition();
            return;
        }

        ByteBuffer bb = ByteBuffer.wrap(segmentIterator.key());
        long sstart = bb.getLong(RdbStorageEngine.TBS_INDEX_SIZE);

        while (true) {
            boolean beyondStop = addRecords(segmentIterator.key(), segmentIterator.value());
            if (beyondStop) {
                stopReached = true;
            }

            segmentIterator.next();
            if (!segmentIterator.isValid()) {
                readNextPartition();
                break;
            }
            bb = ByteBuffer.wrap(segmentIterator.key());
            long g = bb.getLong();
            if (g != sstart) {
                break;
            }
        }
    }

    public void close() {
        if (rdb != null) {
            if (segmentIterator != null) {
                segmentIterator.close();
                segmentIterator = null;
            }
            tablespace.dispose(rdb);
            rdb = null;
        }
    }

    // add all records from this segment into the queue
    // if the stop has been reached add only partially the records, return true
    private boolean addRecords(byte[] key, byte[] val) {
        ByteBuffer kbb = ByteBuffer.wrap(key, TBS_INDEX_SIZE, key.length - TBS_INDEX_SIZE);
        long sstart = kbb.getLong();
        byte[] columnv = new byte[kbb.remaining()];
        kbb.get(columnv);
        ByteBuffer vbb = ByteBuffer.wrap(val);
        HistogramRecord r = null;
        while (vbb.hasRemaining()) {
            long start = sstart * HistogramSegment.GROUPING_FACTOR + vbb.getInt();

            long stop = sstart * HistogramSegment.GROUPING_FACTOR + vbb.getInt();
            int num = vbb.getShort();
            if ((interval.hasStart()) && (stop < interval.getStart())) {
                continue;
            }
            if ((interval.hasEnd()) && (start > interval.getEnd())) {
                if (r != null) {
                    records.add(r);
                }
                return true;
            }
            if (r == null) {
                r = new HistogramRecord(columnv, start, stop, num);
            } else {
                if (start - r.getStop() < mergeTime) {
                    r = new HistogramRecord(r.getColumnv(), r.getStart(), stop, r.getNumTuples() + num);
                } else {
                    records.add(r);
                    r = new HistogramRecord(columnv, start, stop, num);
                }
            }
        }
        if (r != null) {
            records.add(r);
        }
        return false;
{'fraction_non_alphanumeric': 0.09978360182736234, 'fraction_numerical': 0.0019235393123346958, 'mean_word_length': 2.2143225141679546, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5109300', 'n_tokens_mistral': 3525, 'n_tokens_neox': 3410, 'n_words': 771}