library/cpp/actors/util/affinity.h


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

#pragma once

#include "defs.h"
#include "cpumask.h"

// Platform-specific class to set or get thread affinity
class TAffinity: public TThrRefBase, TNonCopyable {
    class TImpl;
    THolder<TImpl> Impl;

public:
    TAffinity();
    TAffinity(const ui8* cpus, ui32 size);
    explicit TAffinity(const TCpuMask& mask);
    ~TAffinity();

    void Current();
    void Set() const;
    bool Empty() const;

    operator TCpuMask() const;
};

// Scoped affinity setter
class TAffinityGuard : TNonCopyable {
    bool Stacked;
    TAffinity OldAffinity;

public:
    TAffinityGuard(const TAffinity* affinity) {
        Stacked = false; 
        if (affinity && !affinity->Empty()) {
            OldAffinity.Current();
            affinity->Set();
            Stacked = true;
        }
    }

    ~TAffinityGuard() {
        Release();
    }

    void Release() {
        if (Stacked) {
            OldAffinity.Set();
            Stacked = false;
        }
    }
};