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