github.com/Rookout/GoSDK@v0.1.48/pkg/aug_manager/aug_manager.go (about)

     1  package aug_manager
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/Rookout/GoSDK/pkg/augs/locations"
     7  	"github.com/Rookout/GoSDK/pkg/com_ws"
     8  	"github.com/Rookout/GoSDK/pkg/logger"
     9  	"github.com/Rookout/GoSDK/pkg/processor"
    10  	"github.com/Rookout/GoSDK/pkg/services/instrumentation"
    11  	"github.com/Rookout/GoSDK/pkg/types"
    12  )
    13  
    14  type AugManager interface {
    15  	InitializeAugs(map[types.AugID]types.AugConfiguration)
    16  	AddAug(types.AugConfiguration)
    17  	RemoveAug(types.AugID) error
    18  	ClearAugs()
    19  }
    20  
    21  type augManager struct {
    22  	triggerServices *instrumentation.TriggerServices
    23  	output          com_ws.Output
    24  	factory         *locations.LocationFactory
    25  
    26  	augIDs     map[types.AugID]interface{}
    27  	augIDsLock sync.Mutex
    28  }
    29  
    30  func NewAugManager(triggerServices *instrumentation.TriggerServices, output com_ws.Output) *augManager {
    31  	augFactory := locations.NewLocationFactory(output, processor.NewProcessorFactory())
    32  	return &augManager{triggerServices: triggerServices, output: output, factory: augFactory, augIDs: make(map[types.AugID]interface{}), augIDsLock: sync.Mutex{}}
    33  }
    34  
    35  func (a *augManager) InitializeAugs(augConfigs map[types.AugID]types.AugConfiguration) {
    36  	a.augIDsLock.Lock()
    37  	defer a.augIDsLock.Unlock()
    38  
    39  	leftovers := make(map[types.AugID]struct{})
    40  	for k := range a.augIDs {
    41  		leftovers[k] = struct{}{}
    42  	}
    43  
    44  	for augID, augConf := range augConfigs {
    45  		if _, ok := leftovers[augID]; ok {
    46  			delete(leftovers, augID)
    47  		} else {
    48  			a.addAug(augConf)
    49  		}
    50  	}
    51  
    52  	for augID := range leftovers {
    53  		err := a.removeAug(augID)
    54  		if err != nil {
    55  			logger.Logger().WithError(err).Errorf("failed to remove leftover aug (%s)", augID)
    56  		}
    57  	}
    58  }
    59  
    60  func (a *augManager) addAug(configuration types.AugConfiguration) {
    61  	aug, err := a.factory.GetLocation(configuration)
    62  	if err != nil {
    63  		logger.Logger().WithError(err).Errorln("Failed to parse aug")
    64  
    65  		if augID, ok := configuration["id"].(types.AugID); ok {
    66  			
    67  			_ = a.output.SendRuleStatus(augID, "Error", err)
    68  		}
    69  		return
    70  	}
    71  
    72  	if _, exists := a.augIDs[aug.GetAugID()]; exists {
    73  		logger.Logger().Debugf("Aug already exists - %s\n", aug.GetAugID())
    74  		return
    75  	}
    76  
    77  	a.triggerServices.GetInstrumentation().AddAug(aug)
    78  	a.augIDs[aug.GetAugID()] = struct{}{}
    79  }
    80  
    81  func (a *augManager) AddAug(configuration types.AugConfiguration) {
    82  	a.augIDsLock.Lock()
    83  	defer a.augIDsLock.Unlock()
    84  
    85  	a.addAug(configuration)
    86  }
    87  
    88  func (a *augManager) removeAug(augID types.AugID) error {
    89  	logger.Logger().Debugf("Removing aug - %s\n", augID)
    90  
    91  	err := a.triggerServices.RemoveAug(augID)
    92  	if err != nil {
    93  		return err
    94  	}
    95  
    96  	delete(a.augIDs, augID)
    97  	return nil
    98  }
    99  
   100  func (a *augManager) RemoveAug(augID types.AugID) error {
   101  	a.augIDsLock.Lock()
   102  	defer a.augIDsLock.Unlock()
   103  
   104  	return a.removeAug(augID)
   105  }
   106  
   107  func (a *augManager) ClearAugs() {
   108  	logger.Logger().Debugf("Clearing all augs\n")
   109  
   110  	var idsCopy []types.AugID
   111  	for k := range a.augIDs {
   112  		idsCopy = append(idsCopy, k)
   113  	}
   114  
   115  	for _, augID := range idsCopy {
   116  		err := a.RemoveAug(augID)
   117  		if err != nil {
   118  			logger.Logger().WithError(err).Errorf("failed to clear aug (%s)", augID)
   119  		}
   120  	}
   121  
   122  	a.triggerServices.ClearAugs()
   123  }