github.com/apptainer/singularity@v3.1.1+incompatible/internal/pkg/cgroups/cgroups_test.go (about)

     1  // Copyright (c) 2018, Sylabs Inc. All rights reserved.
     2  // This software is licensed under a 3-clause BSD license. Please consult the
     3  // LICENSE.md file distributed with the sources of this project regarding your
     4  // rights to use or distribute this software.
     5  
     6  package cgroups
     7  
     8  import (
     9  	"bufio"
    10  	"fmt"
    11  	"io/ioutil"
    12  	"os"
    13  	"os/exec"
    14  	"path/filepath"
    15  	"strconv"
    16  	"strings"
    17  	"testing"
    18  
    19  	"github.com/sylabs/singularity/internal/pkg/test"
    20  )
    21  
    22  func readIntFromFile(path string) (int64, error) {
    23  	file, err := os.Open(path)
    24  	if err != nil {
    25  		return 0, err
    26  	}
    27  	defer file.Close()
    28  
    29  	scanner := bufio.NewScanner(file)
    30  	for scanner.Scan() {
    31  		return strconv.ParseInt(scanner.Text(), 10, 64)
    32  	}
    33  
    34  	return 0, fmt.Errorf("no data found")
    35  }
    36  
    37  func TestCgroups(t *testing.T) {
    38  	test.EnsurePrivilege(t)
    39  
    40  	cmd := exec.Command("/bin/cat")
    41  	pipe, err := cmd.StdinPipe()
    42  	if err != nil {
    43  		t.Fatal(err)
    44  	}
    45  
    46  	if err := cmd.Start(); err != nil {
    47  		t.Fatal(err)
    48  	}
    49  
    50  	pid := cmd.Process.Pid
    51  	strPid := strconv.Itoa(pid)
    52  	path := filepath.Join("/singularity", strPid)
    53  
    54  	manager := &Manager{Pid: pid, Path: path}
    55  	if err := manager.ApplyFromFile("example/cgroups.toml"); err != nil {
    56  		t.Errorf("%s", err)
    57  	}
    58  
    59  	defer manager.Remove()
    60  
    61  	rootPath := manager.GetCgroupRootPath()
    62  	if rootPath == "" {
    63  		t.Fatalf("can't determine cgroups root path, is cgroups enabled ?")
    64  	}
    65  
    66  	cpuShares := filepath.Join(rootPath, "cpu", path, "cpu.shares")
    67  
    68  	i, err := readIntFromFile(cpuShares)
    69  	if err != nil {
    70  		t.Errorf("failed to read %s: %s", cpuShares, err)
    71  	}
    72  	if i != 1024 {
    73  		t.Errorf("cpu shares should be equal to 1024")
    74  	}
    75  
    76  	content := []byte("[cpu]\nshares = 512")
    77  	tmpfile, err := ioutil.TempFile("", "cgroups")
    78  	if err != nil {
    79  		t.Fatal(err)
    80  	}
    81  
    82  	defer os.Remove(tmpfile.Name())
    83  
    84  	if _, err := tmpfile.Write(content); err != nil {
    85  		t.Fatal(err)
    86  	}
    87  	if err := tmpfile.Close(); err != nil {
    88  		t.Fatal(err)
    89  	}
    90  
    91  	// test update/load from PID
    92  	manager = &Manager{Pid: pid}
    93  
    94  	if err := manager.UpdateFromFile(tmpfile.Name()); err != nil {
    95  		t.Fatal(err)
    96  	}
    97  	i, err = readIntFromFile(cpuShares)
    98  	if err != nil {
    99  		t.Errorf("failed to read %s: %s", cpuShares, err)
   100  	}
   101  	if i != 512 {
   102  		t.Errorf("cpu shares should be equal to 512")
   103  	}
   104  
   105  	pipe.Close()
   106  
   107  	cmd.Wait()
   108  }
   109  
   110  func TestPauseResume(t *testing.T) {
   111  	test.EnsurePrivilege(t)
   112  
   113  	manager := &Manager{}
   114  	if err := manager.Pause(); err == nil {
   115  		t.Errorf("unexpected success with PID 0")
   116  	}
   117  	if err := manager.Resume(); err == nil {
   118  		t.Errorf("unexpected success with PID 0")
   119  	}
   120  
   121  	cmd := exec.Command("/bin/cat")
   122  	pipe, err := cmd.StdinPipe()
   123  	if err != nil {
   124  		t.Fatal(err)
   125  	}
   126  
   127  	if err := cmd.Start(); err != nil {
   128  		t.Fatal(err)
   129  	}
   130  
   131  	manager.Pid = cmd.Process.Pid
   132  	manager.Path = filepath.Join("/singularity", strconv.Itoa(manager.Pid))
   133  
   134  	if err := manager.ApplyFromFile("example/cgroups.toml"); err != nil {
   135  		t.Errorf("%s", err)
   136  	}
   137  
   138  	manager.Pause()
   139  
   140  	file, err := os.Open(fmt.Sprintf("/proc/%d/status", manager.Pid))
   141  	if err != nil {
   142  		t.Error(err)
   143  	}
   144  
   145  	scanner := bufio.NewScanner(file)
   146  	stateOk := false
   147  
   148  	for scanner.Scan() {
   149  		if strings.HasPrefix(scanner.Text(), "State:\tD") {
   150  			stateOk = true
   151  			break
   152  		}
   153  	}
   154  
   155  	if !stateOk {
   156  		t.Errorf("failed to pause process %d", manager.Pid)
   157  	}
   158  
   159  	file.Close()
   160  
   161  	manager.Resume()
   162  
   163  	file, err = os.Open(fmt.Sprintf("/proc/%d/status", manager.Pid))
   164  	if err != nil {
   165  		t.Error(err)
   166  	}
   167  
   168  	scanner = bufio.NewScanner(file)
   169  	stateOk = false
   170  
   171  	for scanner.Scan() {
   172  		text := scanner.Text()
   173  		if strings.HasPrefix(text, "State:\tS") || strings.HasPrefix(text, "State:\tR") {
   174  			stateOk = true
   175  			break
   176  		}
   177  	}
   178  
   179  	if !stateOk {
   180  		t.Errorf("failed to resume process %d", manager.Pid)
   181  	}
   182  
   183  	file.Close()
   184  
   185  	defer manager.Remove()
   186  
   187  	pipe.Close()
   188  
   189  	cmd.Wait()
   190  }