public/saasapi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

增加命令功能

Browse Source
This commit is contained in:
algotao
2025-04-07 10:42:08 +08:00
parent 0344c09ce7
commit 74b0033e45
18 changed files with 1844 additions and 494 deletions
Download Patch File Download Diff File Expand all files Collapse all files

164
cmd/saastool/convert.go
View File

@@ -2,27 +2,38 @@ package main
import (
"bufio"
"bytes"
"flag"
"fmt"
"os"
"path"
"runtime"
"strings"
"sync"
"e.coding.net/rta/public/saasapi"
"google.golang.org/protobuf/encoding/protojson"
)
// TODO 转换加速
const (
convertBatchSize = 100000
convertedExt = ".converted"
)
type convertParams struct {
targetCfg *TargetConfig
mapCfg *MapConfig
sourcePath string
destPath string
}
type convertResult struct {
resultBuf bytes.Buffer
convertedLines int
}
func RunConvert(args ...string) error {
fs := flag.NewFlagSet("convert", flag.ExitOnError)
targetCfgFile := paramTargets(fs)
mapCfgFile := paramMap(fs)
sourcePath := paramSourcePath(fs)
destPath := paramDestPath(fs)
@@ -31,19 +42,19 @@ func RunConvert(args ...string) error {
return err
}
if fs.NArg() > 0 || *targetCfgFile == "" || len(*sourcePath) == 0 || len(*destPath) == 0 {
if fs.NArg() > 0 || *mapCfgFile == "" || len(*sourcePath) == 0 || len(*destPath) == 0 {
fs.PrintDefaults()
return nil
}
targetCfg, err := LoadTargetFile(*targetCfgFile)
mapCfg, err := LoadMapFile(*mapCfgFile)
if err != nil {
fmt.Println("LoadConfigFile error", "err", err)
return err
}
convertParams := convertParams{
targetCfg: targetCfg,
mapCfg: mapCfg,
sourcePath: *sourcePath,
destPath: *destPath,
}
@@ -97,7 +108,7 @@ func doFileConvert(convertParams convertParams) error {
os.MkdirAll(convertParams.destPath, os.ModePerm)
}
destName := path.Join(convertParams.destPath, path.Base(convertParams.sourcePath)+".converted")
destName := path.Join(convertParams.destPath, path.Base(convertParams.sourcePath)+convertedExt)
destFile, err := os.Create(destName)
if err != nil {
return err
@@ -108,15 +119,92 @@ func doFileConvert(convertParams convertParams) error {
destWriter := bufio.NewWriter(destFile)
defer destWriter.Flush()
jasonMarshal := protojson.MarshalOptions{Multiline: false, Indent: ""}
// 启动处理协程
workers := []chan []string{}
results := []chan convertResult{}
processedLine := 0
wg := sync.WaitGroup{}
convertMaxWorkers := runtime.GOMAXPROCS(0)
for range convertMaxWorkers {
workerChan := make(chan []string)
workers = append(workers, workerChan)
resultChan := make(chan convertResult)
results = append(results, resultChan)
go func(workerChan <-chan []string, resultChan chan<- convertResult) {
for lines := range workerChan {
convertBatch(lines, convertParams, resultChan)
}
}(workerChan, resultChan)
}
// 启动写入协程
go func() {
i := 0
// TIP: 不要改成for range
for {
select {
case result, ok := <-results[i%convertMaxWorkers]:
if !ok {
return
}
destWriter.Write(result.resultBuf.Bytes())
destWriter.Flush()
processedLine += result.convertedLines
fmt.Printf("\rconverted records: %v [%v]", processedLine, destName)
i++
wg.Done()
}
}
}()
// 读取文件并塞给协程处理
batch := []string{}
batchCount := 0
for scaner.Scan() {
line := scaner.Text()
if line == "" {
continue
}
batch = append(batch, line)
if len(batch) == convertBatchSize {
// 将batch写入协程
wg.Add(1)
workers[batchCount%convertMaxWorkers] <- batch
batch = nil
batchCount++
}
}
if len(batch) > 0 {
wg.Add(1)
workers[batchCount%convertMaxWorkers] <- batch
}
wg.Wait()
// 关闭所有工作协程的通道
for _, workerChan := range workers {
close(workerChan)
}
for _, resultChan := range results {
close(resultChan)
}
fmt.Println("")
return nil
}
func convertBatch(lines []string, convertParams convertParams, resultChan chan<- convertResult) {
byteBuf := bytes.Buffer{}
byteBuf.Grow(1024 * 1024 * 10)
jasonMarshal := protojson.MarshalOptions{Multiline: false, Indent: ""}
for _, line := range lines {
// 按\t分割为两列
parts := strings.Split(line, "\t")
if len(parts) != 2 {
@@ -125,65 +213,63 @@ func doFileConvert(convertParams convertParams) error {
// 读取userid
userid := parts[0]
if len(userid) == 0 {
continue
}
value := parts[1]
value = strings.ReplaceAll(value, "[", "")
value = strings.ReplaceAll(value, "]", "")
// 第二列解析为string数组
targets := strings.Split(value, " ")
saasWriteCmd := &saasapi.WriteCmd{
saasWriteItem := &saasapi.WriteItem{
Userid: userid,
}
if len(userid) == 0 || len(targets) == 0 {
continue
}
// 遍历targets，转换成saasapi.WriteCmd
for _, target := range targets {
if targetinfo, ok := convertParams.targetCfg.Targets[target]; ok {
if targetinfo, ok := convertParams.mapCfg.Targets[target]; ok {
if targetinfo.WriteByte != nil {
if saasWriteCmd.WriteBytes == nil {
saasWriteCmd.WriteBytes = &saasapi.Bytes{}
// 转换byte区
if saasWriteItem.WriteBytes == nil {
saasWriteItem.WriteBytes = &saasapi.Bytes{}
}
saasWriteCmd.WriteBytes.Bytes = append(saasWriteCmd.WriteBytes.Bytes, *targetinfo.WriteByte)
saasWriteItem.WriteBytes.Bytes = append(saasWriteItem.WriteBytes.Bytes, *targetinfo.WriteByte)
if targetinfo.WriteBytePos < 64 {
saasWriteCmd.WriteBytes.Index_1 |= 1 << targetinfo.WriteBytePos
saasWriteItem.WriteBytes.Index_1 |= 1 << targetinfo.WriteBytePos
} else if targetinfo.WriteBytePos < 128 {
saasWriteCmd.WriteBytes.Index_2 |= 1 << (targetinfo.WriteBytePos - 64)
saasWriteItem.WriteBytes.Index_2 |= 1 << (targetinfo.WriteBytePos - 64)
}
}
if targetinfo.WriteUint32 != nil {
if saasWriteCmd.WriteUint32S == nil {
saasWriteCmd.WriteUint32S = &saasapi.Uint32S{}
// 转换uint32区
if saasWriteItem.WriteUint32S == nil {
saasWriteItem.WriteUint32S = &saasapi.Uint32S{}
}
saasWriteCmd.WriteUint32S.Uint32S = append(saasWriteCmd.WriteUint32S.Uint32S, *targetinfo.WriteUint32)
saasWriteCmd.WriteUint32S.Index_1 |= 1 << targetinfo.WriteUint32Pos
saasWriteItem.WriteUint32S.Uint32S = append(saasWriteItem.WriteUint32S.Uint32S, *targetinfo.WriteUint32)
saasWriteItem.WriteUint32S.Index_1 |= 1 << targetinfo.WriteUint32Pos
}
if targetinfo.WriteFlag != nil && targetinfo.WriteExpire != nil {
if saasWriteCmd.WriteFlagsWithExpire == nil {
saasWriteCmd.WriteFlagsWithExpire = &saasapi.FlagsWithExpire{}
// 转换flag区
if saasWriteItem.WriteFlagsWithExpire == nil {
saasWriteItem.WriteFlagsWithExpire = &saasapi.FlagsWithExpire{}
}
saasWriteCmd.WriteFlagsWithExpire.FlagsWithExpire = append(
saasWriteCmd.WriteFlagsWithExpire.FlagsWithExpire, &saasapi.FlagWithExpire{
saasWriteItem.WriteFlagsWithExpire.FlagsWithExpire = append(
saasWriteItem.WriteFlagsWithExpire.FlagsWithExpire, &saasapi.FlagWithExpire{
Flag: *targetinfo.WriteFlag,
Expire: *targetinfo.WriteExpire,
})
saasWriteCmd.WriteFlagsWithExpire.Index_1 |= 1 << targetinfo.WriteFlagWithExpirePos
saasWriteItem.WriteFlagsWithExpire.Index_1 |= 1 << targetinfo.WriteFlagWithExpirePos
}
}
}
// 写入文件
destWriter.WriteString(jasonMarshal.Format(saasWriteCmd))
destWriter.WriteString("\n")
processedLine++
if processedLine%100000 == 0 {
fmt.Printf("\rconverted records: %v [%v]", processedLine, destName)
}
byteBuf.WriteString(jasonMarshal.Format(saasWriteItem))
byteBuf.WriteString("\n")
}
fmt.Printf("\rconverted records: %v [%v]\n", processedLine, destName)
return nil
resultChan <- convertResult{byteBuf, len(lines)}
}