162306a36Sopenharmony_ci// SPDX-License-Identifier: GPL-2.0
262306a36Sopenharmony_ci#include "comm.h"
362306a36Sopenharmony_ci#include <errno.h>
462306a36Sopenharmony_ci#include <stdlib.h>
562306a36Sopenharmony_ci#include <stdio.h>
662306a36Sopenharmony_ci#include <string.h>
762306a36Sopenharmony_ci#include <linux/refcount.h>
862306a36Sopenharmony_ci#include <linux/rbtree.h>
962306a36Sopenharmony_ci#include <linux/zalloc.h>
1062306a36Sopenharmony_ci#include "rwsem.h"
1162306a36Sopenharmony_ci
1262306a36Sopenharmony_cistruct comm_str {
1362306a36Sopenharmony_ci	char *str;
1462306a36Sopenharmony_ci	struct rb_node rb_node;
1562306a36Sopenharmony_ci	refcount_t refcnt;
1662306a36Sopenharmony_ci};
1762306a36Sopenharmony_ci
1862306a36Sopenharmony_ci/* Should perhaps be moved to struct machine */
1962306a36Sopenharmony_cistatic struct rb_root comm_str_root;
2062306a36Sopenharmony_cistatic struct rw_semaphore comm_str_lock = {.lock = PTHREAD_RWLOCK_INITIALIZER,};
2162306a36Sopenharmony_ci
2262306a36Sopenharmony_cistatic struct comm_str *comm_str__get(struct comm_str *cs)
2362306a36Sopenharmony_ci{
2462306a36Sopenharmony_ci	if (cs && refcount_inc_not_zero(&cs->refcnt))
2562306a36Sopenharmony_ci		return cs;
2662306a36Sopenharmony_ci
2762306a36Sopenharmony_ci	return NULL;
2862306a36Sopenharmony_ci}
2962306a36Sopenharmony_ci
3062306a36Sopenharmony_cistatic void comm_str__put(struct comm_str *cs)
3162306a36Sopenharmony_ci{
3262306a36Sopenharmony_ci	if (cs && refcount_dec_and_test(&cs->refcnt)) {
3362306a36Sopenharmony_ci		down_write(&comm_str_lock);
3462306a36Sopenharmony_ci		rb_erase(&cs->rb_node, &comm_str_root);
3562306a36Sopenharmony_ci		up_write(&comm_str_lock);
3662306a36Sopenharmony_ci		zfree(&cs->str);
3762306a36Sopenharmony_ci		free(cs);
3862306a36Sopenharmony_ci	}
3962306a36Sopenharmony_ci}
4062306a36Sopenharmony_ci
4162306a36Sopenharmony_cistatic struct comm_str *comm_str__alloc(const char *str)
4262306a36Sopenharmony_ci{
4362306a36Sopenharmony_ci	struct comm_str *cs;
4462306a36Sopenharmony_ci
4562306a36Sopenharmony_ci	cs = zalloc(sizeof(*cs));
4662306a36Sopenharmony_ci	if (!cs)
4762306a36Sopenharmony_ci		return NULL;
4862306a36Sopenharmony_ci
4962306a36Sopenharmony_ci	cs->str = strdup(str);
5062306a36Sopenharmony_ci	if (!cs->str) {
5162306a36Sopenharmony_ci		free(cs);
5262306a36Sopenharmony_ci		return NULL;
5362306a36Sopenharmony_ci	}
5462306a36Sopenharmony_ci
5562306a36Sopenharmony_ci	refcount_set(&cs->refcnt, 1);
5662306a36Sopenharmony_ci
5762306a36Sopenharmony_ci	return cs;
5862306a36Sopenharmony_ci}
5962306a36Sopenharmony_ci
6062306a36Sopenharmony_cistatic
6162306a36Sopenharmony_cistruct comm_str *__comm_str__findnew(const char *str, struct rb_root *root)
6262306a36Sopenharmony_ci{
6362306a36Sopenharmony_ci	struct rb_node **p = &root->rb_node;
6462306a36Sopenharmony_ci	struct rb_node *parent = NULL;
6562306a36Sopenharmony_ci	struct comm_str *iter, *new;
6662306a36Sopenharmony_ci	int cmp;
6762306a36Sopenharmony_ci
6862306a36Sopenharmony_ci	while (*p != NULL) {
6962306a36Sopenharmony_ci		parent = *p;
7062306a36Sopenharmony_ci		iter = rb_entry(parent, struct comm_str, rb_node);
7162306a36Sopenharmony_ci
7262306a36Sopenharmony_ci		/*
7362306a36Sopenharmony_ci		 * If we race with comm_str__put, iter->refcnt is 0
7462306a36Sopenharmony_ci		 * and it will be removed within comm_str__put call
7562306a36Sopenharmony_ci		 * shortly, ignore it in this search.
7662306a36Sopenharmony_ci		 */
7762306a36Sopenharmony_ci		cmp = strcmp(str, iter->str);
7862306a36Sopenharmony_ci		if (!cmp && comm_str__get(iter))
7962306a36Sopenharmony_ci			return iter;
8062306a36Sopenharmony_ci
8162306a36Sopenharmony_ci		if (cmp < 0)
8262306a36Sopenharmony_ci			p = &(*p)->rb_left;
8362306a36Sopenharmony_ci		else
8462306a36Sopenharmony_ci			p = &(*p)->rb_right;
8562306a36Sopenharmony_ci	}
8662306a36Sopenharmony_ci
8762306a36Sopenharmony_ci	new = comm_str__alloc(str);
8862306a36Sopenharmony_ci	if (!new)
8962306a36Sopenharmony_ci		return NULL;
9062306a36Sopenharmony_ci
9162306a36Sopenharmony_ci	rb_link_node(&new->rb_node, parent, p);
9262306a36Sopenharmony_ci	rb_insert_color(&new->rb_node, root);
9362306a36Sopenharmony_ci
9462306a36Sopenharmony_ci	return new;
9562306a36Sopenharmony_ci}
9662306a36Sopenharmony_ci
9762306a36Sopenharmony_cistatic struct comm_str *comm_str__findnew(const char *str, struct rb_root *root)
9862306a36Sopenharmony_ci{
9962306a36Sopenharmony_ci	struct comm_str *cs;
10062306a36Sopenharmony_ci
10162306a36Sopenharmony_ci	down_write(&comm_str_lock);
10262306a36Sopenharmony_ci	cs = __comm_str__findnew(str, root);
10362306a36Sopenharmony_ci	up_write(&comm_str_lock);
10462306a36Sopenharmony_ci
10562306a36Sopenharmony_ci	return cs;
10662306a36Sopenharmony_ci}
10762306a36Sopenharmony_ci
10862306a36Sopenharmony_cistruct comm *comm__new(const char *str, u64 timestamp, bool exec)
10962306a36Sopenharmony_ci{
11062306a36Sopenharmony_ci	struct comm *comm = zalloc(sizeof(*comm));
11162306a36Sopenharmony_ci
11262306a36Sopenharmony_ci	if (!comm)
11362306a36Sopenharmony_ci		return NULL;
11462306a36Sopenharmony_ci
11562306a36Sopenharmony_ci	comm->start = timestamp;
11662306a36Sopenharmony_ci	comm->exec = exec;
11762306a36Sopenharmony_ci
11862306a36Sopenharmony_ci	comm->comm_str = comm_str__findnew(str, &comm_str_root);
11962306a36Sopenharmony_ci	if (!comm->comm_str) {
12062306a36Sopenharmony_ci		free(comm);
12162306a36Sopenharmony_ci		return NULL;
12262306a36Sopenharmony_ci	}
12362306a36Sopenharmony_ci
12462306a36Sopenharmony_ci	return comm;
12562306a36Sopenharmony_ci}
12662306a36Sopenharmony_ci
12762306a36Sopenharmony_ciint comm__override(struct comm *comm, const char *str, u64 timestamp, bool exec)
12862306a36Sopenharmony_ci{
12962306a36Sopenharmony_ci	struct comm_str *new, *old = comm->comm_str;
13062306a36Sopenharmony_ci
13162306a36Sopenharmony_ci	new = comm_str__findnew(str, &comm_str_root);
13262306a36Sopenharmony_ci	if (!new)
13362306a36Sopenharmony_ci		return -ENOMEM;
13462306a36Sopenharmony_ci
13562306a36Sopenharmony_ci	comm_str__put(old);
13662306a36Sopenharmony_ci	comm->comm_str = new;
13762306a36Sopenharmony_ci	comm->start = timestamp;
13862306a36Sopenharmony_ci	if (exec)
13962306a36Sopenharmony_ci		comm->exec = true;
14062306a36Sopenharmony_ci
14162306a36Sopenharmony_ci	return 0;
14262306a36Sopenharmony_ci}
14362306a36Sopenharmony_ci
14462306a36Sopenharmony_civoid comm__free(struct comm *comm)
14562306a36Sopenharmony_ci{
14662306a36Sopenharmony_ci	comm_str__put(comm->comm_str);
14762306a36Sopenharmony_ci	free(comm);
14862306a36Sopenharmony_ci}
14962306a36Sopenharmony_ci
15062306a36Sopenharmony_ciconst char *comm__str(const struct comm *comm)
15162306a36Sopenharmony_ci{
15262306a36Sopenharmony_ci	return comm->comm_str->str;
15362306a36Sopenharmony_ci}
154