refactor(as400): reorganize logservice package and improve log processing logic

Author: yllada

as400/logservice/processor.go

@@ -1,8 +1,9 @@
-package agent
+package logservice
 
 import (
 	"context"
 	"errors"
+	"os"
 	"strconv"
 	"strings"
 	"sync"
@@ -11,7 +12,6 @@ import (
 	"github.com/google/uuid"
 	"github.com/threatwinds/go-sdk/plugins"
 
-<<<<<<<< HEAD:as400/logservice/processor.go
 	"github.com/utmstack/UTMStack/as400/agent"
 	"github.com/utmstack/UTMStack/as400/config"
 	"github.com/utmstack/UTMStack/as400/conn"
@@ -21,13 +21,6 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
-========
-	"github.com/utmstack/UTMStack/agent/config"
-	"github.com/utmstack/UTMStack/agent/database"
-	"github.com/utmstack/UTMStack/agent/models"
-	"github.com/utmstack/UTMStack/agent/utils"
-	"github.com/utmstack/UTMStack/shared/fs"
->>>>>>>> origin/v11:agent/agent/logprocessor.go
 )
 
 type LogProcessor struct {
@@ -38,48 +31,31 @@ type LogProcessor struct {
 }
 
 var (
-	processor        LogProcessor
-	processorOnce    sync.Once
-	processorInitErr error
-	LogQueue         = make(chan *plugins.Log, 10000)
-	timeCLeanLogs    = 10 * time.Minute
-
-	// ErrAgentUninstalled is returned when the agent uninstalls itself due to invalid key
-	ErrAgentUninstalled = errors.New("agent uninstalled due to invalid key")
+	processor     LogProcessor
+	processorOnce sync.Once
+	LogQueue      = make(chan *plugins.Log)
+	timeToSleep   = 10 * time.Second
+	timeCLeanLogs = 10 * time.Minute
 )
 
-func GetLogProcessor() (*LogProcessor, error) {
+func GetLogProcessor() LogProcessor {
 	processorOnce.Do(func() {
-		db, err := database.GetDB()
-		if err != nil {
-			processorInitErr = err
-			return
-		}
 		processor = LogProcessor{
-			db:             db,
+			db:             database.GetDB(),
 			connErrWritten: false,
 			ackErrWritten:  false,
 			sendErrWritten: false,
 		}
 	})
-	if processorInitErr != nil {
-		return nil, processorInitErr
-	}
-	return &processor, nil
+	return processor
 }
 
 func (l *LogProcessor) ProcessLogs(cnf *config.Config, ctx context.Context) {
 	go l.CleanCountedLogs()
 
 	for {
-		select {
-		case <-ctx.Done():
-			utils.Logger.Info("ProcessLogs stopping due to context cancellation")
-			return
-		default:
-		}
-
-		connection, err := GetCorrelationConnection(cnf)
+		ctxEof, cancelEof := context.WithCancel(context.Background())
+		connection, err := conn.GetCorrelationConnection(cnf)
 		if err != nil {
 			if !l.connErrWritten {
 				utils.Logger.ErrorF("error connecting to Correlation: %v", err)
@@ -90,27 +66,9 @@ func (l *LogProcessor) ProcessLogs(cnf *config.Config, ctx context.Context) {
 		}
 
 		client := plugins.NewIntegrationClient(connection)
-<<<<<<<< HEAD:as400/logservice/processor.go
 		plClient := createClient(client, ctx, cnf)
-========
-		plClient, err := createClient(client, ctx)
-		if err != nil {
-			if errors.Is(err, ErrAgentUninstalled) {
-				utils.Logger.Info("Agent uninstalled, stopping log processor")
-				return
-			}
-			if errors.Is(err, context.Canceled) {
-				utils.Logger.Info("ProcessLogs stopping due to context cancellation")
-				return
-			}
-			utils.Logger.ErrorF("error creating client: %v", err)
-			continue
-		}
->>>>>>>> origin/v11:agent/agent/logprocessor.go
 		l.connErrWritten = false
 
-		// Create context only after successful client creation to avoid leaks
-		ctxEof, cancelEof := context.WithCancel(context.Background())
 		go l.handleAcknowledgements(plClient, ctxEof, cancelEof)
 		l.processLogs(plClient, ctxEof, cancelEof)
 	}
@@ -124,20 +82,38 @@ func (l *LogProcessor) handleAcknowledgements(plClient plugins.Integration_Proce
 		default:
 			ack, err := plClient.Recv()
 			if err != nil {
-				action := HandleGRPCStreamError(err, "failed to receive ack", &l.ackErrWritten)
-				if action == ActionReconnect {
+				if strings.Contains(err.Error(), "EOF") {
+					time.Sleep(timeToSleep)
 					cancel()
 					return
 				}
-				continue
+				st, ok := status.FromError(err)
+				if ok && (st.Code() == codes.Unavailable || st.Code() == codes.Canceled) {
+					if !l.ackErrWritten {
+						utils.Logger.ErrorF("failed to receive ack: %v", err)
+						l.ackErrWritten = true
+					}
+					time.Sleep(timeToSleep)
+					cancel()
+					return
+				} else {
+					if !l.ackErrWritten {
+						utils.Logger.ErrorF("failed to receive ack: %v", err)
+						l.ackErrWritten = true
+					}
+					time.Sleep(timeToSleep)
+					continue
+				}
 			}
 
 			l.ackErrWritten = false
 
+			l.db.Lock()
 			err = l.db.Update(&models.Log{}, "id", ack.LastId, "processed", true)
 			if err != nil {
 				utils.Logger.ErrorF("failed to update log: %v", err)
 			}
+			l.db.Unlock()
 		}
 	}
 }
@@ -149,28 +125,44 @@ func (l *LogProcessor) processLogs(plClient plugins.Integration_ProcessLogClient
 			utils.Logger.Info("context done, exiting processLogs")
 			return
 		case newLog := <-LogQueue:
-			if newLog.Id == "" {
-				id, err := uuid.NewRandom()
-				if err != nil {
-					utils.Logger.ErrorF("failed to generate uuid: %v", err)
-					continue
-				}
+			id, err := uuid.NewRandom()
+			if err != nil {
+				utils.Logger.ErrorF("failed to generate uuid: %v", err)
+				continue
+			}
 
-				newLog.Id = id.String()
-				err = l.db.Create(&models.Log{ID: newLog.Id, Log: newLog.Raw, Type: newLog.DataType, CreatedAt: time.Now(), DataSource: newLog.DataSource, Processed: false})
-				if err != nil {
-					utils.Logger.ErrorF("failed to save log: %v :log: %s", err, newLog.Raw)
-				}
+			newLog.Id = id.String()
+			l.db.Lock()
+			err = l.db.Create(&models.Log{ID: newLog.Id, Log: newLog.Raw, Type: newLog.DataType, CreatedAt: time.Now(), DataSource: newLog.DataSource, Processed: false})
+			if err != nil {
+				utils.Logger.ErrorF("failed to save log: %v :log: %s", err, newLog.Raw)
 			}
+			l.db.Unlock()
 
-			err := plClient.Send(newLog)
+			err = plClient.Send(newLog)
 			if err != nil {
-				action := HandleGRPCStreamError(err, "failed to send log", &l.sendErrWritten)
-				if action == ActionReconnect {
+				if strings.Contains(err.Error(), "EOF") {
+					time.Sleep(timeToSleep)
 					cancel()
 					return
 				}
-				continue
+				st, ok := status.FromError(err)
+				if ok && (st.Code() == codes.Unavailable || st.Code() == codes.Canceled) {
+					if !l.sendErrWritten {
+						utils.Logger.ErrorF("failed to send log: %v :log: %s", err, newLog.Raw)
+						l.sendErrWritten = true
+					}
+					time.Sleep(timeToSleep)
+					cancel()
+					return
+				} else {
+					if !l.sendErrWritten {
+						utils.Logger.ErrorF("failed to send log: %v :log: %s", err, newLog.Raw)
+						l.sendErrWritten = true
+					}
+					time.Sleep(timeToSleep)
+					continue
+				}
 			}
 			l.sendErrWritten = false
 		}
@@ -186,13 +178,17 @@ func (l *LogProcessor) CleanCountedLogs() {
 			utils.Logger.ErrorF("error getting data retention: %s", err)
 			continue
 		}
+		l.db.Lock()
 		_, err = l.db.DeleteOld(&models.Log{}, dataRetention)
 		if err != nil {
 			utils.Logger.ErrorF("error deleting old logs: %s", err)
 		}
+		l.db.Unlock()
 
 		unprocessed := make([]models.Log, 0, 10)
+		l.db.Lock()
 		found, err := l.db.Find(&unprocessed, "processed", false)
+		l.db.Unlock()
 		if err != nil {
 			utils.Logger.ErrorF("error finding unprocessed logs: %s", err)
 			continue
@@ -212,18 +208,10 @@ func (l *LogProcessor) CleanCountedLogs() {
 	}
 }
 
-<<<<<<<< HEAD:as400/logservice/processor.go
 func createClient(client plugins.IntegrationClient, ctx context.Context, cnf *config.Config) plugins.Integration_ProcessLogClient {
-========
-func createClient(client plugins.IntegrationClient, ctx context.Context) (plugins.Integration_ProcessLogClient, error) {
->>>>>>>> origin/v11:agent/agent/logprocessor.go
 	var connErrMsgWritten bool
 	invalidKeyCounter := 0
-	invalidKeyDelay := timeToSleep
-	maxInvalidKeyDelay := 5 * time.Minute
-	maxInvalidKeyAttempts := 100 // ~8+ hours with backoff before uninstall
 	for {
-<<<<<<<< HEAD:as400/logservice/processor.go
 		authCtx := metadata.AppendToOutgoingContext(ctx,
 			"key", cnf.CollectorKey,
 			"id", strconv.Itoa(int(cnf.CollectorID)),
@@ -237,30 +225,9 @@ func createClient(client plugins.IntegrationClient, ctx context.Context) (plugin
 					utils.Logger.Info("Uninstalling collector: reason: collector has been removed from the panel...")
 					_ = agent.UninstallAll()
 					os.Exit(1)
-========
-		select {
-		case <-ctx.Done():
-			return nil, ctx.Err()
-		default:
-		}
-
-		plClient, err := client.ProcessLog(ctx)
-		if err != nil {
-			if strings.Contains(err.Error(), "invalid agent key") {
-				invalidKeyCounter++
-				utils.Logger.ErrorF("invalid agent key (attempt %d/%d), retrying in %v", invalidKeyCounter, maxInvalidKeyAttempts, invalidKeyDelay)
-				if invalidKeyCounter >= maxInvalidKeyAttempts {
-					utils.Logger.ErrorF("uninstalling agent after %d consecutive invalid key errors", maxInvalidKeyAttempts)
-					_ = UninstallAll()
-					return nil, ErrAgentUninstalled
->>>>>>>> origin/v11:agent/agent/logprocessor.go
 				}
-				time.Sleep(invalidKeyDelay)
-				invalidKeyDelay = utils.IncrementReconnectDelay(invalidKeyDelay, maxInvalidKeyDelay)
-				continue
 			} else {
 				invalidKeyCounter = 0
-				invalidKeyDelay = timeToSleep
 			}
 			if !connErrMsgWritten {
 				utils.Logger.ErrorF("failed to create input client: %v", err)
@@ -269,7 +236,7 @@ func createClient(client plugins.IntegrationClient, ctx context.Context) (plugin
 			time.Sleep(timeToSleep)
 			continue
 		}
-		return plClient, nil
+		return plClient
 	}
 }
 
@@ -287,12 +254,12 @@ func SetDataRetention(retention string) error {
 		return errors.New("retention must be greater than 0")
 	}
 
-	return fs.WriteJSON(config.RetentionConfigFile, models.DataRetention{Retention: retentionInt})
+	return utils.WriteJSON(config.RetentionConfigFile, models.DataRetention{Retention: retentionInt})
 }
 
 func GetDataRetention() (int, error) {
 	retention := models.DataRetention{}
-	err := fs.ReadJSON(config.RetentionConfigFile, &retention)
+	err := utils.ReadJson(config.RetentionConfigFile, &retention)
 	if err != nil {
 		return 0, err
 	}