75 lines
1.9 KiB
Go
75 lines
1.9 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"flag"
|
|
"sync"
|
|
|
|
"wave_record/comtrade"
|
|
"wave_record/config"
|
|
"wave_record/database"
|
|
"wave_record/log"
|
|
|
|
"go.mongodb.org/mongo-driver/mongo"
|
|
"go.uber.org/zap"
|
|
|
|
"github.com/panjf2000/ants/v2"
|
|
)
|
|
|
|
var (
|
|
comtradeConfigDir = flag.String("comtrade_config_dir", "./config", "config file dir of wave record project")
|
|
comtradeConfigName = flag.String("comtrade_config_name", "config", "config file name of wave record project")
|
|
comtradeConfigType = flag.String("comtrade_config_type", "yaml", "config file type of wave record project")
|
|
)
|
|
|
|
var (
|
|
waveRecordConfig config.WaveRecordConfig
|
|
addChan chan string
|
|
delChan chan string
|
|
comtradeMap sync.Map
|
|
mongoDBClient *mongo.Client
|
|
logger *zap.Logger
|
|
)
|
|
|
|
func init() {
|
|
addChan = make(chan string, 10)
|
|
delChan = make(chan string, 10)
|
|
}
|
|
|
|
func main() {
|
|
flag.Parse()
|
|
ctx := context.TODO()
|
|
|
|
waveRecordConfig = config.ReadAndInitConfig(*comtradeConfigDir, *comtradeConfigName, *comtradeConfigType)
|
|
|
|
// init mongoDBClient
|
|
mongoDBClient = database.GetMongoDBInstance(ctx, waveRecordConfig.MongoDBURI)
|
|
|
|
// init logger
|
|
logger = log.GetLoggerInstance(waveRecordConfig.LCfg)
|
|
|
|
defer func() {
|
|
if err := mongoDBClient.Disconnect(ctx); err != nil {
|
|
panic(err)
|
|
}
|
|
}()
|
|
|
|
defer logger.Sync()
|
|
|
|
defer ants.Release()
|
|
|
|
pool, err := ants.NewPoolWithFunc(waveRecordConfig.ParseConcurrentQuantity, comtrade.ParseFunc)
|
|
if err != nil {
|
|
logger.Error("init concurrent parse task pool failed", zap.Error(err))
|
|
panic(err)
|
|
}
|
|
|
|
go comtrade.TraverseMonitorDir(ctx, waveRecordConfig.MonitorDir, waveRecordConfig.MongoDBDataBase, &comtradeMap, addChan, delChan, pool)
|
|
|
|
go comtrade.DirWatch(waveRecordConfig.MonitorDir, addChan)
|
|
|
|
go comtrade.MoveComtradeFile(waveRecordConfig.BackupDir, &comtradeMap, delChan)
|
|
|
|
comtrade.ParseComtradeFile(ctx, waveRecordConfig.MonitorDir, waveRecordConfig.MongoDBDataBase, addChan, delChan, &comtradeMap, pool)
|
|
}
|