batch encode opts

This commit is contained in:
thloyi
2026-04-16 17:56:17 +08:00
parent d2879efcc6
commit 1dd843c393
2 changed files with 168 additions and 12 deletions

View File

@@ -1,6 +1,7 @@
package pump_parser
import (
"bufio"
"bytes"
"encoding/binary"
"fmt"
@@ -107,6 +108,18 @@ type TxsBinaryReaderSource interface {
OpenTxsBinaryReader() (io.ReadCloser, error)
}
type TxsBinaryBatchHeaderContext struct {
SourceIndex int
BatchIndex int
Reader *bufio.Reader
}
type TxsBinaryBatchHeaderFunc func(ctx *TxsBinaryBatchHeaderContext) (skip bool, err error)
type TxsBinaryMergeOptions struct {
BatchHeaderFunc TxsBinaryBatchHeaderFunc
}
type PlatformBinary struct {
Platform string
PlatformFee uint64
@@ -307,24 +320,32 @@ func DecodeTxsBinaryReader(r io.Reader) iter.Seq2[*Tx, error] {
}
func MergeTxsBinaryBytes(encodedBatches [][]byte) ([]byte, error) {
return MergeTxsBinaryBytesWithOptions(encodedBatches, TxsBinaryMergeOptions{})
}
func MergeTxsBinaryBytesWithOptions(encodedBatches [][]byte, opts TxsBinaryMergeOptions) ([]byte, error) {
sources := make([]TxsBinaryReaderSource, 0, len(encodedBatches))
for _, encoded := range encodedBatches {
sources = append(sources, txBinaryBytesSource{data: encoded})
}
var out bytes.Buffer
if err := MergeTxsBinarySourcesToWriter(sources, &out); err != nil {
if err := MergeTxsBinarySourcesToWriterWithOptions(sources, &out, opts); err != nil {
return nil, err
}
return out.Bytes(), nil
}
func MergeTxsBinarySourcesToWriter(sources []TxsBinaryReaderSource, w io.Writer) error {
return MergeTxsBinarySourcesToWriterWithOptions(sources, w, TxsBinaryMergeOptions{})
}
func MergeTxsBinarySourcesToWriterWithOptions(sources []TxsBinaryReaderSource, w io.Writer, opts TxsBinaryMergeOptions) error {
if w == nil {
return fmt.Errorf("txs binary writer is nil")
}
plan, err := txBinaryBuildMergePlan(sources)
plan, err := txBinaryBuildMergePlan(sources, opts)
if err != nil {
return err
}
@@ -343,9 +364,22 @@ func MergeTxsBinarySourcesToWriter(sources []TxsBinaryReaderSource, w io.Writer)
return fmt.Errorf("source[%d]: open reader: %w", sourceIndex, err)
}
dec := txBinaryStreamDecoder{reader: reader}
bufferedReader := bufio.NewReader(reader)
dec := txBinaryStreamDecoder{reader: bufferedReader}
batchIndex := 0
for {
skipBatch, err := txBinaryApplyMergeBatchHeader(bufferedReader, opts, sourceIndex, batchIndex)
if err != nil {
closeErr := reader.Close()
if err == io.EOF {
if closeErr != nil {
return fmt.Errorf("source[%d]: close reader: %w", sourceIndex, closeErr)
}
break
}
return fmt.Errorf("source[%d].batch[%d]: %w", sourceIndex, batchIndex, err)
}
header, err := dec.readTxsBinaryHeaderOrEOF()
if err != nil {
closeErr := reader.Close()
@@ -368,6 +402,9 @@ func MergeTxsBinarySourcesToWriter(sources []TxsBinaryReaderSource, w io.Writer)
reader.Close()
return fmt.Errorf("source[%d].batch[%d].tx[%d]: %w", sourceIndex, batchIndex, txIndex, err)
}
if skipBatch {
continue
}
if err := txBinaryRemapTxAddressTable(&tx, header.addressTable, plan.addressTable, plan.addressIndex); err != nil {
reader.Close()
return fmt.Errorf("source[%d].batch[%d].tx[%d]: %w", sourceIndex, batchIndex, txIndex, err)
@@ -1780,7 +1817,7 @@ func txBinaryReadTxBody(dec txBinaryBodyReader, tx *TxBinary, enumTable *txBinar
return nil
}
func txBinaryBuildMergePlan(sources []TxsBinaryReaderSource) (*txsBinaryMergePlan, error) {
func txBinaryBuildMergePlan(sources []TxsBinaryReaderSource, opts TxsBinaryMergeOptions) (*txsBinaryMergePlan, error) {
if len(sources) == 0 {
return nil, fmt.Errorf("txs binary sources are empty")
}
@@ -1801,9 +1838,22 @@ func txBinaryBuildMergePlan(sources []TxsBinaryReaderSource) (*txsBinaryMergePla
return nil, fmt.Errorf("source[%d]: open reader: %w", sourceIndex, err)
}
dec := txBinaryStreamDecoder{reader: reader}
bufferedReader := bufio.NewReader(reader)
dec := txBinaryStreamDecoder{reader: bufferedReader}
batchIndex := 0
for {
skipBatch, err := txBinaryApplyMergeBatchHeader(bufferedReader, opts, sourceIndex, batchIndex)
if err != nil {
closeErr := reader.Close()
if err == io.EOF {
if closeErr != nil {
return nil, fmt.Errorf("source[%d]: close reader: %w", sourceIndex, closeErr)
}
break
}
return nil, fmt.Errorf("source[%d].batch[%d]: %w", sourceIndex, batchIndex, err)
}
header, err := dec.readTxsBinaryHeaderOrEOF()
if err != nil {
closeErr := reader.Close()
@@ -1833,17 +1883,21 @@ func txBinaryBuildMergePlan(sources []TxsBinaryReaderSource) (*txsBinaryMergePla
}
for addressIndex, address := range header.addressTable {
if err := builder.add(address); err != nil {
reader.Close()
return nil, fmt.Errorf("source[%d].batch[%d].address[%d]: %w", sourceIndex, batchIndex, addressIndex, err)
if !skipBatch {
if err := builder.add(address); err != nil {
reader.Close()
return nil, fmt.Errorf("source[%d].batch[%d].address[%d]: %w", sourceIndex, batchIndex, addressIndex, err)
}
}
}
if uint64(plan.txCount)+uint64(header.count) > uint64(math.MaxUint32) {
reader.Close()
return nil, fmt.Errorf("merged tx count exceeds uint32 capacity")
if !skipBatch {
if uint64(plan.txCount)+uint64(header.count) > uint64(math.MaxUint32) {
reader.Close()
return nil, fmt.Errorf("merged tx count exceeds uint32 capacity")
}
plan.txCount += header.count
}
plan.txCount += header.count
for txIndex := uint32(0); txIndex < header.count; txIndex++ {
tx := TxBinary{
@@ -1947,6 +2001,17 @@ func txBinaryWriteAll(w io.Writer, data []byte) error {
return nil
}
func txBinaryApplyMergeBatchHeader(reader *bufio.Reader, opts TxsBinaryMergeOptions, sourceIndex int, batchIndex int) (bool, error) {
if opts.BatchHeaderFunc == nil {
return false, nil
}
return opts.BatchHeaderFunc(&TxsBinaryBatchHeaderContext{
SourceIndex: sourceIndex,
BatchIndex: batchIndex,
Reader: reader,
})
}
type txBinaryEnumTable struct {
version uint16
programs txBinaryEnumSet