1 // SPDX-License-Identifier: GPL-2.0 2 #define _GNU_SOURCE 3 4 #include <errno.h> 5 #include <linux/genetlink.h> 6 #include <linux/netlink.h> 7 #include <linux/taskstats.h> 8 #include <pthread.h> 9 #include <stdbool.h> 10 #include <stdint.h> 11 #include <stdio.h> 12 #include <stdlib.h> 13 #include <string.h> 14 #include <sys/socket.h> 15 #include <sys/types.h> 16 #include <time.h> 17 #include <unistd.h> 18 19 #include "kselftest.h" 20 21 #ifndef NLA_ALIGN 22 #define NLA_ALIGNTO 4 23 #define NLA_ALIGN(len) (((len) + NLA_ALIGNTO - 1) & ~(NLA_ALIGNTO - 1)) 24 #define NLA_HDRLEN ((int)NLA_ALIGN(sizeof(struct nlattr))) 25 #endif 26 27 #define BUSY_NS (200ULL * 1000 * 1000) 28 29 struct worker_ctx { 30 pthread_mutex_t lock; 31 pthread_cond_t cond; 32 bool ready; 33 bool release; 34 }; 35 36 static unsigned long busy_sink; 37 38 static void *taskstats_nla_data(const struct nlattr *na) 39 { 40 return (void *)((char *)na + NLA_HDRLEN); 41 } 42 43 static bool taskstats_nla_ok(const struct nlattr *na, int remaining) 44 { 45 return remaining >= (int)sizeof(*na) && 46 na->nla_len >= sizeof(*na) && 47 na->nla_len <= remaining; 48 } 49 50 static struct nlattr *taskstats_nla_next(const struct nlattr *na, int *remaining) 51 { 52 int aligned_len = NLA_ALIGN(na->nla_len); 53 54 *remaining -= aligned_len; 55 return (struct nlattr *)((char *)na + aligned_len); 56 } 57 58 static uint64_t timespec_diff_ns(const struct timespec *start, 59 const struct timespec *end) 60 { 61 return (uint64_t)(end->tv_sec - start->tv_sec) * 1000000000ULL + 62 (uint64_t)(end->tv_nsec - start->tv_nsec); 63 } 64 65 static void burn_cpu_for_ns(uint64_t runtime_ns) 66 { 67 struct timespec start, now; 68 unsigned long acc = 0; 69 70 if (clock_gettime(CLOCK_MONOTONIC, &start)) { 71 perror("clock_gettime"); 72 exit(EXIT_FAILURE); 73 } 74 75 do { 76 for (int i = 0; i < 100000; i++) 77 acc += i; 78 if (clock_gettime(CLOCK_MONOTONIC, &now)) { 79 perror("clock_gettime"); 80 exit(EXIT_FAILURE); 81 } 82 } while (timespec_diff_ns(&start, &now) < runtime_ns); 83 84 busy_sink = acc; 85 } 86 87 static int netlink_open(void) 88 { 89 struct sockaddr_nl addr = { 90 .nl_family = AF_NETLINK, 91 .nl_pid = getpid(), 92 }; 93 int fd; 94 95 fd = socket(AF_NETLINK, SOCK_RAW, NETLINK_GENERIC); 96 if (fd < 0) 97 return -errno; 98 99 if (bind(fd, (struct sockaddr *)&addr, sizeof(addr)) < 0) { 100 int err = -errno; 101 102 close(fd); 103 return err; 104 } 105 106 return fd; 107 } 108 109 static int send_request(int fd, void *buf, size_t len) 110 { 111 struct sockaddr_nl addr = { 112 .nl_family = AF_NETLINK, 113 }; 114 115 if (sendto(fd, buf, len, 0, (struct sockaddr *)&addr, sizeof(addr)) < 0) 116 return -errno; 117 118 return 0; 119 } 120 121 static int get_family_id(int fd, const char *name) 122 { 123 struct { 124 struct nlmsghdr nlh; 125 struct genlmsghdr genl; 126 char buf[256]; 127 } req = { 0 }; 128 char resp[8192]; 129 struct nlmsghdr *nlh; 130 struct genlmsghdr *genl; 131 struct nlattr *na; 132 int len; 133 int rem; 134 int ret; 135 136 req.nlh.nlmsg_len = NLMSG_LENGTH(GENL_HDRLEN); 137 req.nlh.nlmsg_type = GENL_ID_CTRL; 138 req.nlh.nlmsg_flags = NLM_F_REQUEST; 139 req.nlh.nlmsg_seq = 1; 140 req.nlh.nlmsg_pid = getpid(); 141 142 req.genl.cmd = CTRL_CMD_GETFAMILY; 143 req.genl.version = 1; 144 145 na = (struct nlattr *)((char *)&req + NLMSG_ALIGN(req.nlh.nlmsg_len)); 146 na->nla_type = CTRL_ATTR_FAMILY_NAME; 147 na->nla_len = NLA_HDRLEN + strlen(name) + 1; 148 memcpy(taskstats_nla_data(na), name, strlen(name) + 1); 149 req.nlh.nlmsg_len = NLMSG_ALIGN(req.nlh.nlmsg_len) + NLA_ALIGN(na->nla_len); 150 151 ret = send_request(fd, &req, req.nlh.nlmsg_len); 152 if (ret) 153 return ret; 154 155 len = recv(fd, resp, sizeof(resp), 0); 156 if (len < 0) 157 return -errno; 158 159 for (nlh = (struct nlmsghdr *)resp; NLMSG_OK(nlh, len); 160 nlh = NLMSG_NEXT(nlh, len)) { 161 if (nlh->nlmsg_type == NLMSG_ERROR) { 162 struct nlmsgerr *err = NLMSG_DATA(nlh); 163 164 return err->error ? err->error : -ENOENT; 165 } 166 167 genl = (struct genlmsghdr *)NLMSG_DATA(nlh); 168 rem = nlh->nlmsg_len - NLMSG_HDRLEN - GENL_HDRLEN; 169 na = (struct nlattr *)((char *)genl + GENL_HDRLEN); 170 while (taskstats_nla_ok(na, rem)) { 171 if (na->nla_type == CTRL_ATTR_FAMILY_ID) 172 return *(uint16_t *)taskstats_nla_data(na); 173 na = taskstats_nla_next(na, &rem); 174 } 175 } 176 177 return -ENOENT; 178 } 179 180 static int get_taskstats(int fd, int family_id, uint16_t attr_type, uint32_t id, 181 struct taskstats *stats) 182 { 183 struct { 184 struct nlmsghdr nlh; 185 struct genlmsghdr genl; 186 char buf[256]; 187 } req = { 0 }; 188 char resp[16384]; 189 struct nlmsghdr *nlh; 190 struct genlmsghdr *genl; 191 struct nlattr *na; 192 struct nlattr *nested; 193 int len; 194 int rem; 195 int nrem; 196 int ret; 197 198 memset(stats, 0, sizeof(*stats)); 199 200 req.nlh.nlmsg_len = NLMSG_LENGTH(GENL_HDRLEN); 201 req.nlh.nlmsg_type = family_id; 202 req.nlh.nlmsg_flags = NLM_F_REQUEST; 203 req.nlh.nlmsg_seq = 2; 204 req.nlh.nlmsg_pid = getpid(); 205 206 req.genl.cmd = TASKSTATS_CMD_GET; 207 req.genl.version = 1; 208 209 na = (struct nlattr *)((char *)&req + NLMSG_ALIGN(req.nlh.nlmsg_len)); 210 na->nla_type = attr_type; 211 na->nla_len = NLA_HDRLEN + sizeof(id); 212 memcpy(taskstats_nla_data(na), &id, sizeof(id)); 213 req.nlh.nlmsg_len = NLMSG_ALIGN(req.nlh.nlmsg_len) + NLA_ALIGN(na->nla_len); 214 215 ret = send_request(fd, &req, req.nlh.nlmsg_len); 216 if (ret) 217 return ret; 218 219 len = recv(fd, resp, sizeof(resp), 0); 220 if (len < 0) 221 return -errno; 222 223 for (nlh = (struct nlmsghdr *)resp; NLMSG_OK(nlh, len); 224 nlh = NLMSG_NEXT(nlh, len)) { 225 if (nlh->nlmsg_type == NLMSG_ERROR) { 226 struct nlmsgerr *err = NLMSG_DATA(nlh); 227 228 return err->error ? err->error : -ENOENT; 229 } 230 231 genl = (struct genlmsghdr *)NLMSG_DATA(nlh); 232 rem = nlh->nlmsg_len - NLMSG_HDRLEN - GENL_HDRLEN; 233 na = (struct nlattr *)((char *)genl + GENL_HDRLEN); 234 while (taskstats_nla_ok(na, rem)) { 235 if (na->nla_type == TASKSTATS_TYPE_AGGR_PID || 236 na->nla_type == TASKSTATS_TYPE_AGGR_TGID) { 237 nested = (struct nlattr *)taskstats_nla_data(na); 238 nrem = na->nla_len - NLA_HDRLEN; 239 while (taskstats_nla_ok(nested, nrem)) { 240 if (nested->nla_type == TASKSTATS_TYPE_STATS) { 241 memcpy(stats, taskstats_nla_data(nested), 242 sizeof(*stats)); 243 return 0; 244 } 245 nested = taskstats_nla_next(nested, &nrem); 246 } 247 } 248 na = taskstats_nla_next(na, &rem); 249 } 250 } 251 252 return -ENOENT; 253 } 254 255 static uint64_t cpu_total(const struct taskstats *stats) 256 { 257 return (uint64_t)stats->ac_utime + (uint64_t)stats->ac_stime; 258 } 259 260 static void print_stats(const char *label, const struct taskstats *stats) 261 { 262 ksft_print_msg("%s: cpu_total=%llu nvcsw=%llu nivcsw=%llu\n", 263 label, (unsigned long long)cpu_total(stats), 264 (unsigned long long)stats->nvcsw, 265 (unsigned long long)stats->nivcsw); 266 } 267 268 static void *worker_thread(void *arg) 269 { 270 struct worker_ctx *ctx = arg; 271 272 burn_cpu_for_ns(BUSY_NS); 273 274 pthread_mutex_lock(&ctx->lock); 275 ctx->ready = true; 276 pthread_cond_broadcast(&ctx->cond); 277 while (!ctx->release) 278 pthread_cond_wait(&ctx->cond, &ctx->lock); 279 pthread_mutex_unlock(&ctx->lock); 280 281 return NULL; 282 } 283 284 int main(void) 285 { 286 struct worker_ctx ctx = { 287 .lock = PTHREAD_MUTEX_INITIALIZER, 288 .cond = PTHREAD_COND_INITIALIZER, 289 }; 290 struct taskstats before, after; 291 pthread_t thread; 292 pid_t tgid = getpid(); 293 int family_id; 294 int fd; 295 int ret; 296 297 ksft_print_header(); 298 ksft_set_plan(1); 299 300 if (geteuid()) 301 ksft_exit_skip("taskstats_fill_stats_tgid needs root\n"); 302 303 fd = netlink_open(); 304 if (fd < 0) 305 ksft_exit_skip("failed to open generic netlink socket: %s\n", 306 strerror(-fd)); 307 308 family_id = get_family_id(fd, TASKSTATS_GENL_NAME); 309 if (family_id < 0) 310 ksft_exit_skip("taskstats generic netlink family unavailable: %s\n", 311 strerror(-family_id)); 312 313 /* Create worker thread that burns 200ms of CPU */ 314 if (pthread_create(&thread, NULL, worker_thread, &ctx) != 0) 315 ksft_exit_fail_msg("pthread_create failed: %s\n", strerror(errno)); 316 317 /* Wait for worker to finish generating activity */ 318 pthread_mutex_lock(&ctx.lock); 319 while (!ctx.ready) 320 pthread_cond_wait(&ctx.cond, &ctx.lock); 321 pthread_mutex_unlock(&ctx.lock); 322 323 /* 324 * Snapshot A: TGID stats while worker is alive and sleeping. 325 * Contains main thread + worker contributions. 326 */ 327 ret = get_taskstats(fd, family_id, TASKSTATS_CMD_ATTR_TGID, tgid, &before); 328 if (ret) 329 ksft_exit_fail_msg("TGID query before exit failed: %s\n", 330 strerror(-ret)); 331 332 /* Release worker so it can exit, then join (deterministic wait). 333 * 334 * Kernel exit path ordering guarantees: 335 * do_exit() 336 * taskstats_exit() -> fill_tgid_exit() (accumulates worker into signal->stats) 337 * exit_notify() (releases the thread) 338 * do_task_dead() -> __schedule() (wakes joiner) 339 * 340 * So pthread_join() returns only after fill_tgid_exit() has completed. 341 */ 342 pthread_mutex_lock(&ctx.lock); 343 ctx.release = true; 344 pthread_cond_broadcast(&ctx.cond); 345 pthread_mutex_unlock(&ctx.lock); 346 347 pthread_join(thread, NULL); 348 349 /* 350 * Snapshot B: TGID stats after worker has exited. 351 * fill_stats_for_tgid() does: 352 * memcpy(signal->stats) <- includes fill_tgid_exit accumulation 353 * + scan live threads <- only main thread now 354 */ 355 ret = get_taskstats(fd, family_id, TASKSTATS_CMD_ATTR_TGID, tgid, &after); 356 if (ret) 357 ksft_exit_fail_msg("TGID query after exit failed: %s\n", 358 strerror(-ret)); 359 360 print_stats("TGID before worker exit", &before); 361 print_stats("TGID after worker exit", &after); 362 363 /* 364 * The worker burned 200ms of CPU before the first snapshot. 365 * If the kernel correctly retained its contribution via 366 * fill_tgid_exit(), then the TGID CPU total after exit must be at 367 * least as large as the TGID CPU total before exit. 368 */ 369 ksft_test_result(cpu_total(&after) >= cpu_total(&before), 370 "TGID CPU stats should not regress after thread exit\n"); 371 372 close(fd); 373 ksft_finished(); 374 return ksft_get_fail_cnt() ? KSFT_FAIL : KSFT_PASS; 375 } 376