1 // Copyright (C) 2023 The Android Open Source Project
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //      http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 #include <ditto/multithreading_utils.h>
16 
17 #include <sched.h>
18 #include <unistd.h>
19 
20 namespace dittosuite {
21 
IsSet() const22 bool SchedAttr::IsSet() const { return initialized_; }
23 
Set() const24 void SchedAttr::Set() const {
25   if (!initialized_) {
26     LOGF("Setting uninitialized scheduling attributes");
27   }
28 
29   LOGD("Setting scheduling policy [" + std::to_string(sched_attr_.sched_policy) +
30        "] to thread: " + std::to_string(syscall_.GetTid()));
31 
32   int ret = syscall_.SchedSetattr(0 /* self */, sched_attr_, 0 /* still not implemented */);
33   if (ret) {
34     PLOGF("Failed setting scheduling attributes \n" + to_string(sched_attr_) + "\n");
35   }
36 }
37 
operator =(const dittosuiteproto::SchedAttr & pb)38 SchedAttr& SchedAttr::operator=(const dittosuiteproto::SchedAttr& pb) {
39   typedef dittosuiteproto::SchedAttr::AttributesCase SchedType;
40 
41   sched_attr_.size = sizeof(sched_attr_);
42 
43   sched_attr_.sched_flags = pb.flags();
44 
45   switch (pb.attributes_case()) {
46     case SchedType::kOther:
47       switch (pb.other().policy()) {
48         case dittosuiteproto::SchedAttr::SchedOther::OTHER:
49           sched_attr_.sched_policy = SchedPolicy::SchedNormal;
50           break;
51         case dittosuiteproto::SchedAttr::SchedOther::BATCH:
52           sched_attr_.sched_policy = SchedPolicy::SchedBatch;
53           break;
54       }
55       sched_attr_.sched_nice = pb.other().nice();
56       break;
57     case SchedType::kRt:
58       switch (pb.rt().policy()) {
59         case dittosuiteproto::SchedAttr::SchedRt::FIFO:
60           sched_attr_.sched_policy = SchedPolicy::SchedFifo;
61           break;
62         case dittosuiteproto::SchedAttr::SchedRt::RR:
63           sched_attr_.sched_policy = SchedPolicy::SchedRr;
64           break;
65       }
66       if (pb.rt().priority() < 1 || pb.rt().priority() > 99) {
67         LOGF("Scheduling priority should be in the range [1, 99]");
68       }
69       sched_attr_.sched_priority = pb.rt().priority();
70       break;
71     case SchedType::kDeadline:
72       sched_attr_.sched_policy = SchedPolicy::SchedDeadline;
73       sched_attr_.sched_runtime = pb.deadline().runtime();
74       sched_attr_.sched_deadline = pb.deadline().deadline();
75       sched_attr_.sched_period = pb.deadline().period();
76       break;
77     case SchedType::ATTRIBUTES_NOT_SET:
78       LOGF("Missing scheduling attribute");
79       break;
80   }
81 
82   initialized_ = true;
83 
84   return *this;
85 }
86 
Set() const87 void SchedAffinity::Set() const {
88   if (!initialized_) {
89     LOGF("Setting uninitialized affinity attributes");
90   }
91 
92   LOGD("Setting affinity mask [" + std::to_string(mask_) +
93        "] to thread: " + std::to_string(syscall_.GetTid()));
94 
95   cpu_set_t mask;
96   CPU_ZERO(&mask);
97   uint64_t tmp_bitset = mask_;
98   for (unsigned int i=0; i<sizeof(mask_) * 8; ++i) {
99     if (tmp_bitset & 1) {
100       LOGD("Enabling on CPU: " + std::to_string(i));
101       CPU_SET(i, &mask);
102     }
103     tmp_bitset >>= 1;
104   }
105 
106   int ret = sched_setaffinity(0 /* self */, sizeof(mask), &mask);
107   if (ret) {
108     PLOGF("Failed setting scheduling affinity");
109   }
110 }
111 
IsSet() const112 bool SchedAffinity::IsSet() const {
113   return initialized_;
114 }
115 
operator =(const uint64_t mask)116 SchedAffinity& SchedAffinity::operator=(const uint64_t mask) {
117   if (mask == 0) {
118     LOGF("Empty CPU affinity mask");
119   }
120 
121   mask_ = mask;
122   initialized_ = true;
123 
124   return *this;
125 }
126 
127 /*
128  * Create a copy of environ and return the new argv[0] size
129  *
130  * The stack looks pretty much as follows:
131  *
132  * | [...]
133  * |
134  * | argc
135  * | argv[0] (pointer)
136  * | argv[1] (pointer)
137  * | ...
138  * | NULL
139  * | environ[0] (pointer)
140  * | environ[1] (pointer)
141  * | ...
142  * | NULL
143  * |
144  * | [...]
145  * |
146  * | *argv[0] (string)
147  * | *argv[1] (string)
148  * | ...
149  * | *environ[0] (string)
150  * | *environ[1] (string)
151  * | ...
152  *
153  * After this function is call, all the *environ[*] strings will be moved in
154  * dynamic memory, and the old environ strings space in the stack can be used
155  * to extend *argv[*].
156  */
move_environ(int argc,char ** argv)157 int move_environ(int argc, char**argv) {
158   // Count the number of items in environ
159   int env_last_id = -1;
160   if (environ) {
161     while (environ[++env_last_id])
162       ;
163   }
164 
165   unsigned int argv_strings_size;
166   if (env_last_id > 0) {
167     // If there is something in environ (it exists and there is something more
168     // than the terminating NULL), size is the total size that will be usable
169     // for argv after environ is moved.  In fact, this would be the size of all
170     // the argv strings, plus the size of all the current environ strings.
171     // More specifically, this is:
172     // - the address of the last element of environ,
173     // - plus its content size, so now we have the very last byte of environ,
174     // - subtracted the address of the first string of argv.
175     argv_strings_size = environ[env_last_id - 1] + strlen(environ[env_last_id - 1]) - argv[0];
176   } else {
177     // Otherwise, this is just the size of all the argv strings.
178     argv_strings_size = argv[argc - 1] + strlen(argv[argc - 1]) - argv[0];
179   }
180 
181   if (environ) {
182     // Create a copy of environ in dynamic memory
183     char** new_environ = static_cast<char**>(malloc(env_last_id * sizeof(char*)));
184 
185     // Create a copy in dynamic memory for all the environ strings
186     unsigned int i = -1;
187     while (environ[++i]) {
188       new_environ[i] = strdup(environ[i]);
189     }
190 
191     // Also, update the environ pointer
192     environ = new_environ;
193   }
194 
195   return argv_strings_size;
196 }
197 
198 /*
199  * Update the title of the calling process by updating argv[0] (may be
200  * destructive for the following argv[1..])
201  *
202  * If the length of the new title is larger than the previously allocated
203  * argv[0] in the stack, then this function will overwrite the next argv
204  * strings.
205  * If the length of the new title is larger than all the argv strings, this
206  * function will move environ and use its reclaimed space.
207  */
setproctitle(int argc,char ** argv,const char * title)208 void setproctitle(int argc, char** argv, const char* title) {
209   size_t available_size = strlen(argv[0]);
210   size_t argv_size = argv[argc - 1] + strlen(argv[argc - 1]) - argv[0];
211 
212   if (available_size < argv_size) {
213     available_size = move_environ(argc, argv);
214   }
215 
216   memset(argv[0], 0, available_size);
217   snprintf(argv[0], available_size - 1, "%s", title);
218 }
219 
220 }  // namespace dittosuite
221