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