feat: optimize code (#149)

Co-authored-by: wencaiwulue <895703375@qq.com>

pkg/controlplane/main.go

@@ -10,14 +10,15 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
-func Main(filename string, port uint, logger *log.Logger) {
+func Main(ctx context.Context, filename string, port uint, logger *log.Logger) error {
 	snapshotCache := cache.NewSnapshotCache(false, cache.IDHash{}, logger)
 	proc := NewProcessor(snapshotCache, logger)
 
+	errChan := make(chan error, 2)
+
 	go func() {
-		ctx := context.Background()
 		server := serverv3.NewServer(ctx, snapshotCache, nil)
-		RunServer(ctx, server, port)
+		errChan <- RunServer(ctx, server, port)
 	}()
 
 	notifyCh := make(chan NotifyMessage, 100)
@@ -29,20 +30,29 @@ func Main(filename string, port uint, logger *log.Logger) {
 
 	watcher, err := fsnotify.NewWatcher()
 	if err != nil {
-		log.Fatal(fmt.Errorf("failed to create file watcher, err: %v", err))
+		return fmt.Errorf("failed to create file watcher: %v", err)
 	}
 	defer watcher.Close()
-	if err = watcher.Add(filename); err != nil {
-		log.Fatal(fmt.Errorf("failed to add file: %s to wather, err: %v", filename, err))
+	err = watcher.Add(filename)
+	if err != nil {
+		return fmt.Errorf("failed to add file: %s to wather: %v", filename, err)
 	}
 	go func() {
-		log.Fatal(Watch(watcher, filename, notifyCh))
+		errChan <- Watch(watcher, filename, notifyCh)
 	}()
 
 	for {
 		select {
 		case msg := <-notifyCh:
-			proc.ProcessFile(msg)
+			err = proc.ProcessFile(msg)
+			if err != nil {
+				log.Errorf("failed to process file: %v", err)
+				return err
+			}
+		case err = <-errChan:
+			return err
+		case <-ctx.Done():
+			return ctx.Err()
 		}
 	}
 }

pkg/controlplane/processor.go

@@ -44,11 +44,11 @@ func (p *Processor) newVersion() string {
 	return strconv.FormatInt(p.version, 10)
 }
 
-func (p *Processor) ProcessFile(file NotifyMessage) {
+func (p *Processor) ProcessFile(file NotifyMessage) error {
 	configList, err := ParseYaml(file.FilePath)
 	if err != nil {
 		p.logger.Errorf("error parsing yaml file: %+v", err)
-		return
+		return err
 	}
 	for _, config := range configList {
 		if len(config.Uid) == 0 {
@@ -76,21 +76,22 @@ func (p *Processor) ProcessFile(file NotifyMessage) {
 
 		if err != nil {
 			p.logger.Errorf("snapshot inconsistency: %v, err: %v", snapshot, err)
-			return
+			return err
 		}
 
 		if err = snapshot.Consistent(); err != nil {
 			p.logger.Errorf("snapshot inconsistency: %v, err: %v", snapshot, err)
-			return
+			return err
 		}
 		p.logger.Debugf("will serve snapshot %+v, nodeID: %s", snapshot, config.Uid)
 		if err = p.cache.SetSnapshot(context.Background(), config.Uid, snapshot); err != nil {
 			p.logger.Errorf("snapshot error %q for %v", err, snapshot)
-			p.logger.Fatal(err)
+			return err
 		}
 
 		p.expireCache.Set(config.Uid, config, time.Minute*5)
 	}
+	return nil
 }
 
 func ParseYaml(file string) ([]*Virtual, error) {

pkg/controlplane/server.go

@@ -21,13 +21,13 @@ const (
 	grpcMaxConcurrentStreams = 1000000
 )
 
-func RunServer(ctx context.Context, server serverv3.Server, port uint) {
+func RunServer(ctx context.Context, server serverv3.Server, port uint) error {
 	grpcServer := grpc.NewServer(grpc.MaxConcurrentStreams(grpcMaxConcurrentStreams))
 
 	var lc net.ListenConfig
 	listener, err := lc.Listen(ctx, "tcp", fmt.Sprintf(":%d", port))
 	if err != nil {
-		log.Fatal(err)
+		return err
 	}
 
 	discoverygrpc.RegisterAggregatedDiscoveryServiceServer(grpcServer, server)
@@ -39,7 +39,5 @@ func RunServer(ctx context.Context, server serverv3.Server, port uint) {
 	runtimeservice.RegisterRuntimeDiscoveryServiceServer(grpcServer, server)
 
 	log.Infof("management server listening on %d", port)
-	if err = grpcServer.Serve(listener); err != nil {
-		log.Fatal(err)
-	}
+	return grpcServer.Serve(listener)
 }