github.com/castai/kvisor@v1.7.1-0.20240516114728-b3572a2607b5/cmd/event-generator/app/controller.go (about)

     1  package app
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"math/rand"
     7  	"time"
     8  
     9  	"github.com/castai/kvisor/pkg/logging"
    10  	"github.com/samber/lo"
    11  	appsv1 "k8s.io/api/apps/v1"
    12  	corev1 "k8s.io/api/core/v1"
    13  	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
    14  	"k8s.io/client-go/kubernetes"
    15  )
    16  
    17  func newController(log *logging.Logger, client kubernetes.Interface) *controller {
    18  	return &controller{
    19  		log:    log.WithField("component", "ctrl"),
    20  		client: client,
    21  		rnd:    rand.New(rand.NewSource(time.Now().Unix())), //nolint:gosec
    22  	}
    23  }
    24  
    25  type controller struct {
    26  	client kubernetes.Interface
    27  	log    *logging.Logger
    28  
    29  	rnd *rand.Rand
    30  }
    31  
    32  func (c *controller) run(ctx context.Context) error {
    33  	c.log.Info("running kvisor events generator")
    34  	defer c.log.Info("stopping kvisor events generator")
    35  
    36  	go func() {
    37  		for {
    38  			select {
    39  			case <-ctx.Done():
    40  				return
    41  			case <-time.After(5 * time.Minute):
    42  				if err := c.removeRunnerContainer(ctx); err != nil {
    43  					c.log.Errorf("remove runner container: %v", err)
    44  				}
    45  			}
    46  		}
    47  	}()
    48  
    49  	for {
    50  		rndDur := time.Duration(c.rnd.Intn(10)+1) * time.Minute
    51  		c.log.Debugf("next inject after %s", rndDur)
    52  		select {
    53  		case <-ctx.Done():
    54  			return nil
    55  		case <-time.After(rndDur):
    56  			if err := c.injectRunnerContainer(ctx); err != nil {
    57  				c.log.Errorf("inject runner container: %v", err)
    58  			}
    59  		}
    60  	}
    61  }
    62  
    63  func (c *controller) injectRunnerContainer(ctx context.Context) error {
    64  	c.log.Debugf("running inject")
    65  	dep, err := c.selectRandomDeployment(ctx)
    66  	if err != nil {
    67  		return err
    68  	}
    69  	c.log.Debugf("selected deployment, name=%s", dep.Name)
    70  
    71  	if lo.ContainsBy(dep.Spec.Template.Spec.InitContainers, func(item corev1.Container) bool {
    72  		return item.Name == "kvisor-anomaly-runner"
    73  	}) {
    74  		c.log.Debugf("skipping, deployment already contains event runner init container")
    75  		return nil
    76  	}
    77  
    78  	imgName, err := c.getGeneratorImageName(ctx)
    79  	if err != nil {
    80  		return err
    81  	}
    82  
    83  	dep.Spec.Template.Spec.InitContainers = append(dep.Spec.Template.Spec.InitContainers, corev1.Container{
    84  		Name:  "kvisor-anomaly-runner",
    85  		Image: imgName,
    86  		Args: []string{
    87  			"-mode=event-runner",
    88  		},
    89  	})
    90  
    91  	dep.ObjectMeta.Annotations["kvisor-anomaly-runner"] = "true"
    92  	if _, err := c.client.AppsV1().Deployments(dep.Namespace).Update(ctx, dep, metav1.UpdateOptions{}); err != nil {
    93  		return err
    94  	}
    95  
    96  	c.log.Debugf("injected event runner init container")
    97  
    98  	return nil
    99  }
   100  
   101  func (c *controller) removeRunnerContainer(ctx context.Context) error {
   102  	c.log.Debugf("running cleanup")
   103  	list, err := c.client.AppsV1().Deployments(metav1.NamespaceAll).List(ctx, metav1.ListOptions{})
   104  	if err != nil {
   105  		return err
   106  	}
   107  
   108  	for _, item := range list.Items {
   109  		item := item
   110  		if _, found := item.ObjectMeta.Annotations["kvisor-anomaly-runner"]; found {
   111  			item.Spec.Template.Spec.InitContainers = lo.Filter(item.Spec.Template.Spec.InitContainers, func(item corev1.Container, index int) bool {
   112  				return item.Name != "kvisor-anomaly-runner"
   113  			})
   114  			delete(item.ObjectMeta.Annotations, "kvisor-anomaly-runner")
   115  			if _, err := c.client.AppsV1().Deployments(item.Namespace).Update(ctx, &item, metav1.UpdateOptions{}); err != nil {
   116  				return err
   117  			}
   118  			c.log.Debugf("removed event runner init container from deployment, name=%s", item.Name)
   119  		}
   120  	}
   121  	return nil
   122  }
   123  
   124  func (c *controller) selectRandomDeployment(ctx context.Context) (*appsv1.Deployment, error) {
   125  	list, err := c.client.AppsV1().Deployments(metav1.NamespaceAll).List(ctx, metav1.ListOptions{})
   126  	if err != nil {
   127  		return nil, err
   128  	}
   129  	list.Items = lo.Filter(list.Items, func(item appsv1.Deployment, index int) bool {
   130  		return item.Namespace != "kvisor"
   131  	})
   132  
   133  	if len(list.Items) == 0 {
   134  		return nil, errors.New("no deployments")
   135  	}
   136  
   137  	di := c.rnd.Intn(len(list.Items))
   138  	return &list.Items[di], nil
   139  }
   140  
   141  func (c *controller) getGeneratorImageName(ctx context.Context) (string, error) {
   142  	dep, err := c.client.AppsV1().Deployments("kvisor").Get(ctx, "kvisor-event-generator", metav1.GetOptions{})
   143  	if err != nil {
   144  		return "", err
   145  	}
   146  	return dep.Spec.Template.Spec.Containers[0].Image, nil
   147  }