collector.go

Documentation: github.com/chaos-mesh/chaos-mesh/pkg/workflow/task/collector

     1  // Copyright 2021 Chaos Mesh Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package collector
    15  
    16  import (
    17  	"context"
    18  
    19  	"k8s.io/client-go/rest"
    20  	"sigs.k8s.io/controller-runtime/pkg/client"
    21  )
    22  
    23  // Collector is a set of tools for collecting parameters/context from user defined task
    24  type Collector interface {
    25  	CollectContext(ctx context.Context) (env map[string]interface{}, err error)
    26  }
    27  
    28  type ComposeCollector struct {
    29  	collectors []Collector
    30  }
    31  
    32  func (it *ComposeCollector) CollectContext(ctx context.Context) (env map[string]interface{}, err error) {
    33  	if len(it.collectors) == 0 {
    34  		return nil, nil
    35  	}
    36  	if len(it.collectors) == 1 {
    37  		return it.collectors[0].CollectContext(ctx)
    38  	}
    39  
    40  	result := make(map[string]interface{})
    41  	for _, collector := range it.collectors {
    42  		temp, err := collector.CollectContext(ctx)
    43  		if err != nil {
    44  			return nil, err
    45  		}
    46  		mapExtend(result, temp)
    47  	}
    48  	return result, nil
    49  }
    50  
    51  // mapExtend will merge another map into the origin map, value with duplicated key will be replaced.
    52  // origin map should not be nil
    53  func mapExtend(origin map[string]interface{}, another map[string]interface{}) {
    54  	if origin == nil || another == nil {
    55  		return
    56  	}
    57  	for k, v := range another {
    58  		origin[k] = v
    59  	}
    60  }
    61  
    62  func DefaultCollector(kubeClient client.Client, restConfig *rest.Config, namespace, podName, containerName string) Collector {
    63  	return &ComposeCollector{collectors: []Collector{
    64  		NewExitCodeCollector(kubeClient, namespace, podName, containerName),
    65  		NewStdoutCollector(restConfig, namespace, podName, containerName),
    66  	}}
    67  }
    68