551 lines
17 KiB
C++
551 lines
17 KiB
C++
/*
|
||
Description: 排序过程的第二阶段
|
||
* 第二阶段（phase-2）：
|
||
* 1. 打开所有中间文件，并关联上解压线程池
|
||
* 2. phase-2-stage-1：解压中间文件，并进行归并排序，放入buf
|
||
* 3. phase-2-stage-2：多线程进行压缩，放入buf
|
||
* 4. phase-2-stage-3：将buf写入最终的bam文件
|
||
* 5. 要注意计算index并写入index文件
|
||
Copyright : All right reserved by ICT
|
||
|
||
Author : Zhang Zhonghai
|
||
Date : 2026/02/08
|
||
*/
|
||
|
||
#pragma once
|
||
#include <spdlog/spdlog.h>
|
||
#include <stdio.h>
|
||
|
||
#include <string>
|
||
|
||
#include "phase_1.h"
|
||
#include "sam_io.h"
|
||
#include "sort.h"
|
||
using std::string;
|
||
|
||
// 循环缓冲区
|
||
struct CircularBuffer {
|
||
uint8_t* data = nullptr;
|
||
size_t readIdx = 0; // 可以读取的开始位置
|
||
size_t writeIdx = 0; // 可以写入的开始位置
|
||
size_t valueSize = 0; // 有效字节
|
||
size_t bufSize = 0; // 缓冲区空间
|
||
|
||
CircularBuffer() {}
|
||
CircularBuffer(size_t initSize) {
|
||
bufSize = initSize;
|
||
data = (uint8_t*)malloc(bufSize);
|
||
}
|
||
~CircularBuffer() {
|
||
if (data) {
|
||
free(data);
|
||
data = nullptr;
|
||
}
|
||
readIdx = 0;
|
||
writeIdx = 0;
|
||
valueSize = 0;
|
||
bufSize = 0;
|
||
}
|
||
void AllocMem(size_t memSize) { ReAllocMem(memSize); }
|
||
|
||
void ReAllocMem(size_t memSize) {
|
||
if (memSize > bufSize) {
|
||
bufSize = memSize;
|
||
data = (uint8_t*)realloc(data, bufSize);
|
||
}
|
||
}
|
||
|
||
size_t Write(uint8_t* src, size_t len) {
|
||
if (len == 0) return 0;
|
||
size_t freeSpace = bufSize - valueSize;
|
||
size_t toWrite = MIN(len, freeSpace);
|
||
if (toWrite == 0) return 0;
|
||
|
||
size_t firstPart = MIN(toWrite, bufSize - writeIdx);
|
||
// 第一段：从 write_index_ 到末尾
|
||
memcpy(data + writeIdx, src, firstPart);
|
||
// 第二段：如果超过了末尾，则从头部继续
|
||
if (toWrite > firstPart) {
|
||
memcpy(data, src + firstPart, toWrite - firstPart);
|
||
}
|
||
writeIdx = (writeIdx + toWrite) % bufSize;
|
||
valueSize += toWrite;
|
||
return toWrite;
|
||
}
|
||
|
||
size_t Read(uint8_t* out, size_t len) {
|
||
if (len == 0) return 0;
|
||
size_t toRead = MIN(len, valueSize);
|
||
if (toRead == 0) return 0;
|
||
|
||
size_t firstPart = MIN(toRead, bufSize - readIdx);
|
||
memcpy(out, data + readIdx, firstPart);
|
||
if (toRead > firstPart) {
|
||
std::memcpy(out + firstPart, data, toRead - firstPart);
|
||
}
|
||
readIdx = (readIdx + toRead) % bufSize;
|
||
valueSize -= toRead;
|
||
return toRead;
|
||
}
|
||
|
||
// 返回第一个连续空间的大小
|
||
size_t FirstPartWriteSize() {
|
||
size_t freeSpace = bufSize - valueSize;
|
||
size_t firstPart = MIN(freeSpace, bufSize - writeIdx);
|
||
return firstPart;
|
||
}
|
||
|
||
// 如果空间不连续，那么返回第二个连续空间的内存大小
|
||
size_t SecondPartWriteSize() {
|
||
size_t freeSpace = bufSize - valueSize;
|
||
size_t firstPart = MIN(freeSpace, bufSize - writeIdx);
|
||
if (firstPart == freeSpace)
|
||
return 0;
|
||
return freeSpace - firstPart;
|
||
}
|
||
|
||
// 跳过不能完整保存一个bam的空间
|
||
void SkipWrite(size_t skipBytes) {
|
||
valueSize += skipBytes;
|
||
writeIdx = (writeIdx + skipBytes) % bufSize;
|
||
}
|
||
|
||
void SkipRead(size_t from, size_t skipBytes) {
|
||
if (valueSize <= skipBytes) { // 直接清空
|
||
readIdx = 0; // 可以读取的开始位置
|
||
writeIdx = 0; // 可以写入的开始位置
|
||
valueSize = 0; // 有效字节
|
||
} else {
|
||
if (from == readIdx) {
|
||
readIdx = (readIdx + skipBytes) % bufSize;
|
||
valueSize -= skipBytes;
|
||
} else {
|
||
valueSize -= (bufSize - readIdx + skipBytes);
|
||
readIdx = skipBytes;
|
||
}
|
||
}
|
||
}
|
||
|
||
// 退回一个bam
|
||
void RevertRead(size_t from, size_t revertBytes) {
|
||
readIdx = from - revertBytes;
|
||
valueSize += revertBytes;
|
||
}
|
||
|
||
// 读取一个bam
|
||
size_t ReadBam(uint8_t* out, size_t start, size_t len) {
|
||
memcpy(out, data + start, len);
|
||
valueSize -= len;
|
||
if (readIdx != start) {
|
||
valueSize -= bufSize - readIdx;
|
||
}
|
||
readIdx = start + len;
|
||
return len;
|
||
}
|
||
|
||
// 是否为空
|
||
bool Empty() const { return valueSize == 0; }
|
||
// 是否已满
|
||
bool Full() const { return valueSize == bufSize; }
|
||
// 当前数据量（字节）
|
||
size_t Size() const { return valueSize; }
|
||
// 剩余可写空间
|
||
size_t Free() const { return bufSize - valueSize; }
|
||
// 总容量
|
||
size_t Capacity() const { return bufSize; }
|
||
|
||
void Clear() {
|
||
readIdx = 0; // 可以读取的开始位置
|
||
writeIdx = 0; // 可以写入的开始位置
|
||
valueSize = 0; // 有效字节
|
||
}
|
||
};
|
||
|
||
// 循环数组
|
||
template<class T>
|
||
struct CircularArray {
|
||
std::vector<T> arr;
|
||
size_t readIdx = 0; // 可以读取的开始位置
|
||
size_t writeIdx = 0; // 可以写入的开始位置
|
||
size_t valueSize = 0; // 有效字节
|
||
size_t bufSize = 0; // 缓冲区空间
|
||
|
||
CircularArray() {}
|
||
CircularArray(size_t initSize) {
|
||
bufSize = initSize;
|
||
arr.resize(initSize);
|
||
}
|
||
void AllocArr(size_t arrSize) { ReAllocArr(arrSize); }
|
||
|
||
void ReAllocArr(size_t arrSize) {
|
||
if (arrSize > bufSize) {
|
||
bufSize = arrSize;
|
||
arr.resize(arrSize);
|
||
}
|
||
}
|
||
|
||
inline T& operator[](size_t idx) {
|
||
return arr[(idx + readIdx) % bufSize];
|
||
}
|
||
|
||
// 推入元素（自动扩容）
|
||
void Push(const T& value) {
|
||
if (Full()) {
|
||
ReAllocArr(bufSize * 1.5);
|
||
}
|
||
arr[writeIdx] = value;
|
||
writeIdx = (writeIdx + 1) % bufSize;
|
||
++valueSize;
|
||
}
|
||
|
||
void Push(T&& value) {
|
||
if (Full()) {
|
||
ReAllocArr(bufSize * 1.5);
|
||
}
|
||
arr[writeIdx] = std::move(value);
|
||
writeIdx = (writeIdx + 1) % bufSize;
|
||
++valueSize;
|
||
}
|
||
|
||
// 弹出头部元素（如果为空则返回 nullopt）
|
||
void Pop() {
|
||
if (Empty())
|
||
return;
|
||
--valueSize;
|
||
readIdx = (readIdx + 1) % bufSize;
|
||
}
|
||
|
||
void Revert() {
|
||
++valueSize;
|
||
readIdx = (readIdx - 1) % bufSize;
|
||
}
|
||
|
||
// 查看头部/尾部（不弹出）
|
||
T* Front() {
|
||
if (Empty())
|
||
return nullptr;
|
||
return &arr[readIdx];
|
||
}
|
||
|
||
T* Back() {
|
||
if (Empty())
|
||
return nullptr;
|
||
return &arr[(writeIdx - 1) % bufSize];
|
||
}
|
||
|
||
// 是否为空
|
||
bool Empty() const { return valueSize == 0; }
|
||
// 是否已满
|
||
bool Full() const { return valueSize == bufSize; }
|
||
// 当前数据量（字节）
|
||
size_t Size() const { return valueSize; }
|
||
// 剩余可写空间
|
||
size_t Free() const { return bufSize - valueSize; }
|
||
// 总容量
|
||
size_t Capacity() const { return bufSize; }
|
||
|
||
|
||
void Clear() {
|
||
readIdx = 0; // 可以读取的开始位置
|
||
writeIdx = 0; // 可以写入的开始位置
|
||
valueSize = 0; // 有效字节
|
||
}
|
||
};
|
||
|
||
// 合并过程需要的缓冲区
|
||
struct Phase2MergeBuffer {
|
||
CircularBuffer data;
|
||
CircularArray<OneBam> bams;
|
||
bool initialized = false;
|
||
size_t numProcessed = 0;
|
||
|
||
OneBam* Front() { return bams.Front(); }
|
||
void Pop() {
|
||
OneBam* bp = Front();
|
||
data.SkipRead(bp->offset, bp->wholeBamLen);
|
||
bams.Pop();
|
||
}
|
||
|
||
void Revert() {
|
||
bams.Revert();
|
||
OneBam* b = bams.Front();
|
||
data.RevertRead(b->offset, b->wholeBamLen);
|
||
}
|
||
|
||
size_t Size() { return bams.Size(); }
|
||
|
||
bool hasSpace(size_t bamLen) { return data.FirstPartWriteSize() >= bamLen; }
|
||
|
||
// 从一个解压后的block缓冲区拷贝多个bam到循环缓冲区
|
||
// 返回实际拷贝的bam数量
|
||
size_t CopyBams(DataBuffer &blockBuf, BamArr &arr, int start) {
|
||
if (start >= arr.Size()) return 0;
|
||
int origStart = start;
|
||
size_t numCopied = 0;
|
||
OneBam* b1 = &arr.Get(start);
|
||
OneBam* b2 = arr.Back();
|
||
|
||
size_t firstPartSize = data.FirstPartWriteSize();
|
||
size_t secondPartSize = data.SecondPartWriteSize();
|
||
size_t needSize = b2->offset - b1->offset + b2->wholeBamLen;
|
||
int stop = arr.Size();
|
||
if (needSize <= firstPartSize) { // 在第一个连续空间里就能放下
|
||
// 每个bam的offset需要加上diff，以对应新的buf
|
||
int64_t diff = (int64_t)data.writeIdx - b1->offset;
|
||
data.Write(blockBuf.data + b1->offset, needSize);
|
||
for (int i = start; i < stop; ++i) {
|
||
bams.Push(arr.Get(i));
|
||
bams.Back()->offset += diff;
|
||
bams.Back()->addr = data.data;
|
||
}
|
||
} else {
|
||
stop = start;
|
||
size_t firstNeedSize = 0;
|
||
while (arr.Get(stop).wholeBamLen + firstNeedSize < firstPartSize) {
|
||
firstNeedSize += arr.Get(stop).wholeBamLen;
|
||
stop += 1;
|
||
}
|
||
int64_t diff = (int64_t)data.writeIdx - b1->offset;
|
||
data.Write(blockBuf.data + b1->offset, firstNeedSize);
|
||
for (int i = start; i < stop; ++i) {
|
||
bams.Push(arr.Get(i));
|
||
bams.Back()->offset += diff;
|
||
bams.Back()->addr = data.data;
|
||
}
|
||
// 跳过first part不能放下完整bam的部分
|
||
data.SkipWrite(firstPartSize - firstNeedSize);
|
||
b1 = &arr.Get(stop);
|
||
start = stop;
|
||
// 拷贝第二段
|
||
size_t secondNeedSize = needSize - firstNeedSize;
|
||
if (secondNeedSize <= secondPartSize) {
|
||
stop = arr.Size();
|
||
} else {
|
||
secondNeedSize = 0;
|
||
while (arr.Get(stop).wholeBamLen + secondNeedSize < secondPartSize) {
|
||
secondNeedSize += arr.Get(stop).wholeBamLen;
|
||
stop += 1;
|
||
}
|
||
}
|
||
diff = (int64_t)data.writeIdx - b1->offset;
|
||
data.Write(blockBuf.data + b1->offset, secondNeedSize);
|
||
for (int i = start; i < stop; ++i) {
|
||
bams.Push(arr.Get(i));
|
||
bams.Back()->offset += diff;
|
||
bams.Back()->addr = data.data;
|
||
}
|
||
}
|
||
numCopied = stop - origStart;
|
||
|
||
return numCopied;
|
||
}
|
||
|
||
void InitSize(size_t _bufsize, size_t _bamsize) {
|
||
data.AllocMem(_bufsize);
|
||
bams.AllocArr(_bamsize);
|
||
initialized = true;
|
||
}
|
||
};
|
||
|
||
struct BamDataBuffer {
|
||
uint64_t offset = 0;
|
||
int bamOffset = 0;
|
||
DataBuffer block;
|
||
BamArr bams;
|
||
void Clear() {
|
||
offset = 0;
|
||
block.Clear();
|
||
bams.Clear();
|
||
}
|
||
};
|
||
|
||
// 读取和解压需要的buffer
|
||
struct Phase2ReadBuffer : ReadBuffer {
|
||
FastVector<BamDataBuffer> uncompressBlocks;
|
||
int bamNum = 0;
|
||
};
|
||
|
||
// 解析
|
||
struct UncompressBuffer {
|
||
DataBuffer blockBuf; // 解压的block放在这里
|
||
BamArr bamArr; // 解析后的bam数据放在这里
|
||
int startIdx = 0;
|
||
|
||
size_t Size() { return bamArr.Size() - startIdx; }
|
||
|
||
OneBam *Front() {
|
||
if (Size() == 0)
|
||
return nullptr;
|
||
return &bamArr.Get(startIdx);
|
||
}
|
||
|
||
void Clear() {
|
||
startIdx = 0;
|
||
blockBuf.Clear();
|
||
bamArr.Clear();
|
||
}
|
||
};
|
||
|
||
struct Phase2File {
|
||
// static constexpr int MID_FILE_BUF_SIZE = 8 * SINGLE_BLOCK_SIZE;
|
||
static constexpr int READ_BUF_NUM = 2; // 读入的buf数量
|
||
static constexpr int UNCOMPRESSS_BUF_NUM = 2; // 解析的buf数量
|
||
static constexpr int COPY_BUF_NUM = 1; // 拷贝到merge缓冲区数量
|
||
|
||
// 参数
|
||
size_t readBufSize = 0;
|
||
string fileName;
|
||
|
||
FILE* fp = nullptr;
|
||
Phase2ReadBuffer readData[READ_BUF_NUM];
|
||
UncompressBuffer uncompressBuf[UNCOMPRESSS_BUF_NUM];
|
||
DataBuffer halfBlock; // 剩余不完整的压缩的block数据
|
||
|
||
// 读
|
||
volatile int readyReadBufNum = 0; // 有几个buffer已经填满可用了, 类似lock_t的作用
|
||
uint64_t readOrder = 0; // 当前读取的buf索引
|
||
|
||
// 解压
|
||
volatile int readyUncompressBufNum = 0;
|
||
uint64_t uncompressOrder = 0; // 写入的buf索引，类似readOrder
|
||
bool needUncompress = false;
|
||
|
||
// 拷贝到归并buf
|
||
Phase2MergeBuffer mergeData; // 归并时候用到的数据
|
||
uint64_t copyOrder = 0;
|
||
volatile int readyMergeBufNum = 0;
|
||
|
||
volatile int finish = 0; // 读入结束
|
||
|
||
void Init() {
|
||
for (int i = 0; i < READ_BUF_NUM; ++i) {
|
||
readData[i].Resize(readBufSize);
|
||
}
|
||
halfBlock.AllocMem(SINGLE_BLOCK_SIZE);
|
||
// 打开中间文件
|
||
// spdlog::info("file: {}", fileName);
|
||
fp = fopen(fileName.c_str(), "rb");
|
||
}
|
||
|
||
void Init(const string& _fileName, size_t _readBufSize) {
|
||
readBufSize = _readBufSize;
|
||
fileName = _fileName;
|
||
Init();
|
||
}
|
||
|
||
// 缓冲区大小应该根据中间文件的数量进行调整
|
||
|
||
Phase2File() {}
|
||
Phase2File(const string &_fileName, size_t _readBufSize) {
|
||
readBufSize = _readBufSize;
|
||
fileName = _fileName;
|
||
}
|
||
|
||
~Phase2File() {
|
||
if (fp) {
|
||
fclose(fp);
|
||
}
|
||
}
|
||
};
|
||
|
||
// 用于合并压缩的数据结构
|
||
struct Phase2MergeData {
|
||
vector<DataBuffer> blockDataArr; // 待压缩的数据
|
||
vector<DataBuffer> compressDataArr; // 压缩后的数据
|
||
|
||
void Resize(int blockNum) {
|
||
blockDataArr.resize(blockNum);
|
||
compressDataArr.resize(blockNum);
|
||
}
|
||
};
|
||
|
||
/* 第二阶段的多线程流水线参数 */
|
||
struct Phase2PipelineArg {
|
||
// common parameters
|
||
static const int MERGE_BUF_NUM = 2; // 压缩的buf数量
|
||
static const int COMPRESS_BUF_NUM = 2; // 压缩的buf数量
|
||
|
||
int compressBlocksThreshold = 0; // 进行归并排序的，单次block数量阈值，参数
|
||
int numThread = 0; // 线程数
|
||
int compressLevel = 3; // 压缩级别，参数
|
||
uint64_t numBam = 0;
|
||
|
||
// 从phase1转移来的数据
|
||
UncompressBlockBuffer& uncompressData; // 所有线程共用一个，串行往这里添加解压后的block数据
|
||
BamArr& allBams; // 所有线程共用一个，串行往这里添加解析后的bam数据
|
||
|
||
// 读入中间文件
|
||
vector<Phase2File> midFiles;
|
||
yarn::lock_t* readSig;
|
||
uint64_t readOrder = 0;
|
||
volatile int readFinish = 0;
|
||
|
||
// 解压读入的blocks
|
||
yarn::lock_t* uncompressSig;
|
||
uint64_t uncompressOrder = 0;
|
||
volatile int uncompressFinish = 0;
|
||
int curBlockNum = 0;
|
||
|
||
// 拷贝到merge缓冲区
|
||
yarn::lock_t* copyMerge;
|
||
uint64_t copyMergeOrder = 0;
|
||
int uncompressReadyNum = 0;
|
||
volatile int copyMergeFinish = 0;
|
||
|
||
// 将解压的数据拷贝进merge的缓冲区，归并后放入，以便后续压缩用到
|
||
MergeCompressData mergeData[MERGE_BUF_NUM];
|
||
yarn::lock_t* mergeSig;
|
||
yarn::lock_t* copyCompressSig;
|
||
uint64_t mergeOrder = 0; // 前半段
|
||
uint64_t copyCompressOrder = 0; // 后半段，需要压缩
|
||
volatile int mergeFinish = 0;
|
||
volatile int copyCompressFinish = 0;
|
||
|
||
// 压缩
|
||
DataBuffer compressBuf[COMPRESS_BUF_NUM]; // 压缩用到的buf
|
||
yarn::lock_t* compressSig;
|
||
uint64_t compressOrder = 0;
|
||
volatile int compressFinish = 0;
|
||
|
||
// 写入文件
|
||
FILE* outFilePtr = nullptr; // 文件指针
|
||
uint64_t writeOrder = 0;
|
||
|
||
void SetThreadNum(int nThread) {
|
||
numThread = nThread;
|
||
// const int kThres = 1024;
|
||
const int kThres = 2048;
|
||
int threadBlocksThres = numThread * 32;
|
||
compressBlocksThreshold = threadBlocksThres < kThres ? threadBlocksThres : kThres; // 大概64M的解压缩后的数据
|
||
for (int i = 0; i < COMPRESS_BUF_NUM; ++i) {
|
||
mergeData[i].Resize(compressBlocksThreshold);
|
||
}
|
||
}
|
||
|
||
Phase2PipelineArg(UncompressBlockBuffer& lastBuf, BamArr& lastBams, string &midFilePrefix, int midFileNum, int threadNum) : uncompressData(lastBuf), allBams(lastBams){
|
||
SetThreadNum(threadNum);
|
||
|
||
int allFileBufSize = 64L * 1024 * 1024;
|
||
// int bufBlocks = (numThread * 4 + midFileNum) / (midFileNum + 1);
|
||
// size_t kFileBufSize = bufBlocks * SINGLE_BLOCK_SIZE;
|
||
size_t kFileBufSize = allFileBufSize / (midFileNum + 1);
|
||
|
||
spdlog::info("mid file buf size: {}", kFileBufSize);
|
||
|
||
midFiles.resize(midFileNum);
|
||
for (int i = 0; i < midFileNum; ++i) {
|
||
midFiles[i].Init(midFilePrefix + std::to_string(i), kFileBufSize);
|
||
}
|
||
|
||
readSig = yarn::NEW_LOCK(0);
|
||
uncompressSig = yarn::NEW_LOCK(0);
|
||
copyMerge = yarn::NEW_LOCK(0);
|
||
mergeSig = yarn::NEW_LOCK(0);
|
||
copyCompressSig = yarn::NEW_LOCK(0);
|
||
compressSig = yarn::NEW_LOCK(0);
|
||
|
||
|
||
}
|
||
};
|
||
|
||
void phase2Pipeline(Phase2PipelineArg &p); |