pod_container_handler.go

Documentation: github.com/chaos-mesh/chaos-mesh/pkg/chaosdaemon/tasks

     1  // Copyright 2022 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  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package tasks
    16  
    17  import (
    18  	"sync"
    19  
    20  	"github.com/go-logr/logr"
    21  
    22  	"github.com/chaos-mesh/chaos-mesh/pkg/cerr"
    23  )
    24  
    25  var ErrNotPodContainerName = cerr.NotType[PodContainerName]()
    26  
    27  var ErrPodProcessMapNotInit = cerr.NotInit[map[PodContainerName]SysPID]().WrapName("PodContainerNameProcessMap").Err()
    28  
    29  type PodContainerName string
    30  
    31  func (p PodContainerName) ToID() string {
    32  	return string(p)
    33  }
    34  
    35  // ChaosOnPOD stand for the inner process injector for container.
    36  type ChaosOnPOD interface {
    37  	Injectable
    38  	Recoverable
    39  }
    40  
    41  type PodContainerNameProcessMap struct {
    42  	m      map[PodContainerName]SysPID
    43  	rwLock sync.RWMutex
    44  }
    45  
    46  func NewPodProcessMap() PodContainerNameProcessMap {
    47  	return PodContainerNameProcessMap{
    48  		m:      make(map[PodContainerName]SysPID),
    49  		rwLock: sync.RWMutex{},
    50  	}
    51  }
    52  
    53  func (p *PodContainerNameProcessMap) Read(PodContainerName PodContainerName) (SysPID, error) {
    54  	p.rwLock.RLock()
    55  	defer p.rwLock.RUnlock()
    56  	sysPID, ok := p.m[PodContainerName]
    57  	if !ok {
    58  		return SysPID(0), ErrNotFoundSysID.WithStack().Err()
    59  	}
    60  	return sysPID, nil
    61  }
    62  
    63  func (p *PodContainerNameProcessMap) Write(PodContainerName PodContainerName, sysPID SysPID) {
    64  	p.rwLock.Lock()
    65  	defer p.rwLock.Unlock()
    66  	p.m[PodContainerName] = sysPID
    67  }
    68  
    69  func (p *PodContainerNameProcessMap) Delete(podPID PodContainerName) {
    70  	p.rwLock.Lock()
    71  	defer p.rwLock.Unlock()
    72  
    73  	delete(p.m, podPID)
    74  }
    75  
    76  // PodHandler implements injecting & recovering on a kubernetes POD.
    77  type PodHandler struct {
    78  	PodProcessMap *PodContainerNameProcessMap
    79  	SubProcess    ChaosOnPOD
    80  	Logger        logr.Logger
    81  }
    82  
    83  func NewPodHandler(podProcessMap *PodContainerNameProcessMap, sub ChaosOnPOD, logger logr.Logger) PodHandler {
    84  	return PodHandler{
    85  		PodProcessMap: podProcessMap,
    86  		SubProcess:    sub,
    87  		Logger:        logr.New(logger.GetSink()),
    88  	}
    89  }
    90  
    91  // Inject get the container process IsID and Inject it with major injector.
    92  // Be careful about the error handling here.
    93  func (p *PodHandler) Inject(id IsID) error {
    94  	podPID, ok := id.(PodContainerName)
    95  	if !ok {
    96  		return ErrNotPodContainerName.WrapInput(id).Err()
    97  	}
    98  	if p.PodProcessMap == nil {
    99  		return ErrPodProcessMapNotInit
   100  	}
   101  
   102  	sysPID, err := p.PodProcessMap.Read(podPID)
   103  	if err != nil {
   104  		return err
   105  	}
   106  
   107  	err = p.SubProcess.Inject(sysPID)
   108  	return err
   109  }
   110  
   111  // Recover get the container process IsID and Recover it with major injector.
   112  // Be careful about the error handling here.
   113  func (p *PodHandler) Recover(id IsID) error {
   114  	podPID, ok := id.(PodContainerName)
   115  	if !ok {
   116  		return ErrNotPodContainerName.WrapInput(id).Err()
   117  	}
   118  	if p.PodProcessMap == nil {
   119  		return ErrPodProcessMapNotInit
   120  	}
   121  
   122  	sysPID, err := p.PodProcessMap.Read(podPID)
   123  	if err != nil {
   124  		return err
   125  	}
   126  
   127  	err = p.SubProcess.Recover(sysPID)
   128  	return err
   129  }
   130