jy-data-extract/udpdataclear/udpSensitiveWords/grpc_server/main.go

package main

import (
	"context"
	"fmt"
	"github.com/importcjj/sensitive"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"google.golang.org/grpc"
	"gopkg.in/yaml.v3"
	"io/ioutil"
	"log"
	"net"
	"runtime"
	"sensitiveWords.udp/proto_grpc"
	"sensitiveWords.udp/util"
	"strings"
	"time"
)

const (
	PORT     = ":50051"
	YAMLFILE = "./server.yaml"
)

var YamlConfig YAMLConfig
var MixDataMgo *util.MongodbSim
var Filter *sensitive.Filter

func init() {
	yamlFile, err := ioutil.ReadFile(YAMLFILE)
	if err != nil {
		log.Fatalln("load conf error")
	}
	err = yaml.Unmarshal(yamlFile, &YamlConfig)
	if err != nil {
		fmt.Println(err.Error())
	}
	log.Printf("%#v", YamlConfig)
	MixDataMgo = &util.MongodbSim{
		MongodbAddr: YamlConfig.MixdataMgoAddr,
		Size:        YamlConfig.MongodbPoolSize,
		DbName:      YamlConfig.DbName,
		UserName:    YamlConfig.UserName,
		PassWord:    YamlConfig.PassWord,
	}
	MixDataMgo.InitPool()

	gteid, err := primitive.ObjectIDFromHex(YamlConfig.TaskGteId)
	if err != nil {
		log.Fatalln(err)
	}
	lteid, err := primitive.ObjectIDFromHex(YamlConfig.TaskLteId)
	if err != nil {
		log.Fatalln(err)
	}
	iter := MixDataMgo.GetMgoConn().C("unique_qyxy").Find(map[string]interface{}{
		"_id": map[string]interface{}{
			"$gte": gteid,
			"$lte": lteid,
		},
	}).Sort("_id").Iter()
	Filter = sensitive.New()
	var initnum uint
	saveIdArr ,start_id:= make([]map[string]string,0),""
	var m runtime.MemStats
	for tmp := map[string]interface{}{}; iter.Next(&tmp); tmp = map[string]interface{}{} {
		if start_id=="" {
			start_id = BsonTOStringId(tmp["_id"])
		}
		Filter.AddWord(tmp["qy_name"].(string))
		initnum++
		if initnum%10000==0 {
			runtime.ReadMemStats(&m)
			men :=toMegaBytes(m.HeapAlloc)
			log.Printf("current index %d\tos %.2f M",initnum, men)
			if men>500 {
				saveIdArr = append(saveIdArr, map[string]string{
					"start":start_id,
					"end":BsonTOStringId(tmp["_id"]),
				})
				runtime.GC()
				Filter = sensitive.New()
				start_id = ""
				time.Sleep(time.Second*5)
			}
		}
		break

	}

	saveIdArr = append(saveIdArr, map[string]string{
		"start":start_id,
		"end":"",
	})

	for k,v:=range saveIdArr{
		log.Println("第",k,"段",v["start"],v["end"])
	}

	log.Println("init ok", initnum)
}


func toMegaBytes(bytes uint64) float64 {
	return float64(bytes) / 1024 / 1024
}



func main() {
	lis, err := net.Listen("tcp", PORT)
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}
	s := grpc.NewServer()
	proto_grpc.RegisterSensitiveWordsServer(s, &server{})
	log.Println("server start:", PORT)
	if err := s.Serve(lis); err != nil {
		log.Fatalf("failed to serve: %v", err)
	}

}

type server struct {
	proto_grpc.SensitiveWordsServer
}

func (s *server) Search(ctx context.Context, in *proto_grpc.Request) (*proto_grpc.ResultSensitiveWords, error) {
	text := in.GetText()
	log.Println(text)
	findAll := Filter.FindAll(text)
	return &proto_grpc.ResultSensitiveWords{SensitiveWords: strings.Join(findAll, ",")}, nil
}

type YAMLConfig struct {
	MixdataMgoAddr  string `yaml:"mixdataMgoAddr"`
	UserName        string `yaml:"userName"`
	PassWord        string `yaml:"passWord"`
	DbName          string `yaml:"dbName"`
	MongodbPoolSize int    `yaml:"mongodbPoolSize"`
	TaskGteId       string `yaml:"taskGteId"`
	TaskLteId       string `yaml:"taskLteId"`
}
func StringTOBsonId(id string) primitive.ObjectID {
	objectId, _ := primitive.ObjectIDFromHex(id)
	return objectId
}

func BsonTOStringId(id interface{}) string {
	return id.(primitive.ObjectID).Hex()
}