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