1#!/bin/bash
2# SPDX-License-Identifier: GPL-2.0
3
4# Kselftest framework requirement - SKIP code is 4.
5ksft_skip=4
6
7set -e
8
9if [[ $(id -u) -ne 0 ]]; then
10  echo "This test must be run as root. Skipping..."
11  exit $ksft_skip
12fi
13
14usage_file=usage_in_bytes
15
16if [[ "$1" == "-cgroup-v2" ]]; then
17  cgroup2=1
18  usage_file=current
19fi
20
21CGROUP_ROOT='/dev/cgroup/memory'
22MNT='/mnt/huge/'
23
24if [[ ! -e $CGROUP_ROOT ]]; then
25  mkdir -p $CGROUP_ROOT
26  if [[ $cgroup2 ]]; then
27    mount -t cgroup2 none $CGROUP_ROOT
28    sleep 1
29    echo "+hugetlb +memory" >$CGROUP_ROOT/cgroup.subtree_control
30  else
31    mount -t cgroup memory,hugetlb $CGROUP_ROOT
32  fi
33fi
34
35function get_machine_hugepage_size() {
36  hpz=$(grep -i hugepagesize /proc/meminfo)
37  kb=${hpz:14:-3}
38  mb=$(($kb / 1024))
39  echo $mb
40}
41
42MB=$(get_machine_hugepage_size)
43
44function cleanup() {
45  echo cleanup
46  set +e
47  rm -rf "$MNT"/* 2>/dev/null
48  umount "$MNT" 2>/dev/null
49  rmdir "$MNT" 2>/dev/null
50  rmdir "$CGROUP_ROOT"/a/b 2>/dev/null
51  rmdir "$CGROUP_ROOT"/a 2>/dev/null
52  rmdir "$CGROUP_ROOT"/test1 2>/dev/null
53  echo 0 >/proc/sys/vm/nr_hugepages
54  set -e
55}
56
57function assert_state() {
58  local expected_a="$1"
59  local expected_a_hugetlb="$2"
60  local expected_b=""
61  local expected_b_hugetlb=""
62
63  if [ ! -z ${3:-} ] && [ ! -z ${4:-} ]; then
64    expected_b="$3"
65    expected_b_hugetlb="$4"
66  fi
67  local tolerance=$((5 * 1024 * 1024))
68
69  local actual_a
70  actual_a="$(cat "$CGROUP_ROOT"/a/memory.$usage_file)"
71  if [[ $actual_a -lt $(($expected_a - $tolerance)) ]] ||
72    [[ $actual_a -gt $(($expected_a + $tolerance)) ]]; then
73    echo actual a = $((${actual_a%% *} / 1024 / 1024)) MB
74    echo expected a = $((${expected_a%% *} / 1024 / 1024)) MB
75    echo fail
76
77    cleanup
78    exit 1
79  fi
80
81  local actual_a_hugetlb
82  actual_a_hugetlb="$(cat "$CGROUP_ROOT"/a/hugetlb.${MB}MB.$usage_file)"
83  if [[ $actual_a_hugetlb -lt $(($expected_a_hugetlb - $tolerance)) ]] ||
84    [[ $actual_a_hugetlb -gt $(($expected_a_hugetlb + $tolerance)) ]]; then
85    echo actual a hugetlb = $((${actual_a_hugetlb%% *} / 1024 / 1024)) MB
86    echo expected a hugetlb = $((${expected_a_hugetlb%% *} / 1024 / 1024)) MB
87    echo fail
88
89    cleanup
90    exit 1
91  fi
92
93  if [[ -z "$expected_b" || -z "$expected_b_hugetlb" ]]; then
94    return
95  fi
96
97  local actual_b
98  actual_b="$(cat "$CGROUP_ROOT"/a/b/memory.$usage_file)"
99  if [[ $actual_b -lt $(($expected_b - $tolerance)) ]] ||
100    [[ $actual_b -gt $(($expected_b + $tolerance)) ]]; then
101    echo actual b = $((${actual_b%% *} / 1024 / 1024)) MB
102    echo expected b = $((${expected_b%% *} / 1024 / 1024)) MB
103    echo fail
104
105    cleanup
106    exit 1
107  fi
108
109  local actual_b_hugetlb
110  actual_b_hugetlb="$(cat "$CGROUP_ROOT"/a/b/hugetlb.${MB}MB.$usage_file)"
111  if [[ $actual_b_hugetlb -lt $(($expected_b_hugetlb - $tolerance)) ]] ||
112    [[ $actual_b_hugetlb -gt $(($expected_b_hugetlb + $tolerance)) ]]; then
113    echo actual b hugetlb = $((${actual_b_hugetlb%% *} / 1024 / 1024)) MB
114    echo expected b hugetlb = $((${expected_b_hugetlb%% *} / 1024 / 1024)) MB
115    echo fail
116
117    cleanup
118    exit 1
119  fi
120}
121
122function setup() {
123  echo 100 >/proc/sys/vm/nr_hugepages
124  mkdir "$CGROUP_ROOT"/a
125  sleep 1
126  if [[ $cgroup2 ]]; then
127    echo "+hugetlb +memory" >$CGROUP_ROOT/a/cgroup.subtree_control
128  else
129    echo 0 >$CGROUP_ROOT/a/cpuset.mems
130    echo 0 >$CGROUP_ROOT/a/cpuset.cpus
131  fi
132
133  mkdir "$CGROUP_ROOT"/a/b
134
135  if [[ ! $cgroup2 ]]; then
136    echo 0 >$CGROUP_ROOT/a/b/cpuset.mems
137    echo 0 >$CGROUP_ROOT/a/b/cpuset.cpus
138  fi
139
140  mkdir -p "$MNT"
141  mount -t hugetlbfs none "$MNT"
142}
143
144write_hugetlbfs() {
145  local cgroup="$1"
146  local path="$2"
147  local size="$3"
148
149  if [[ $cgroup2 ]]; then
150    echo $$ >$CGROUP_ROOT/$cgroup/cgroup.procs
151  else
152    echo 0 >$CGROUP_ROOT/$cgroup/cpuset.mems
153    echo 0 >$CGROUP_ROOT/$cgroup/cpuset.cpus
154    echo $$ >"$CGROUP_ROOT/$cgroup/tasks"
155  fi
156  ./write_to_hugetlbfs -p "$path" -s "$size" -m 0 -o
157  if [[ $cgroup2 ]]; then
158    echo $$ >$CGROUP_ROOT/cgroup.procs
159  else
160    echo $$ >"$CGROUP_ROOT/tasks"
161  fi
162  echo
163}
164
165set -e
166
167size=$((${MB} * 1024 * 1024 * 25)) # 50MB = 25 * 2MB hugepages.
168
169cleanup
170
171echo
172echo
173echo Test charge, rmdir, uncharge
174setup
175echo mkdir
176mkdir $CGROUP_ROOT/test1
177
178echo write
179write_hugetlbfs test1 "$MNT"/test $size
180
181echo rmdir
182rmdir $CGROUP_ROOT/test1
183mkdir $CGROUP_ROOT/test1
184
185echo uncharge
186rm -rf /mnt/huge/*
187
188cleanup
189
190echo done
191echo
192echo
193if [[ ! $cgroup2 ]]; then
194  echo "Test parent and child hugetlb usage"
195  setup
196
197  echo write
198  write_hugetlbfs a "$MNT"/test $size
199
200  echo Assert memory charged correctly for parent use.
201  assert_state 0 $size 0 0
202
203  write_hugetlbfs a/b "$MNT"/test2 $size
204
205  echo Assert memory charged correctly for child use.
206  assert_state 0 $(($size * 2)) 0 $size
207
208  rmdir "$CGROUP_ROOT"/a/b
209  sleep 5
210  echo Assert memory reparent correctly.
211  assert_state 0 $(($size * 2))
212
213  rm -rf "$MNT"/*
214  umount "$MNT"
215  echo Assert memory uncharged correctly.
216  assert_state 0 0
217
218  cleanup
219fi
220
221echo
222echo
223echo "Test child only hugetlb usage"
224echo setup
225setup
226
227echo write
228write_hugetlbfs a/b "$MNT"/test2 $size
229
230echo Assert memory charged correctly for child only use.
231assert_state 0 $(($size)) 0 $size
232
233rmdir "$CGROUP_ROOT"/a/b
234echo Assert memory reparent correctly.
235assert_state 0 $size
236
237rm -rf "$MNT"/*
238umount "$MNT"
239echo Assert memory uncharged correctly.
240assert_state 0 0
241
242cleanup
243
244echo ALL PASS
245
246umount $CGROUP_ROOT
247rm -rf $CGROUP_ROOT
248