1*d7be555eSGeorge Liu #pragma once
2*d7be555eSGeorge Liu 
3*d7be555eSGeorge Liu #include "NVMeSensor.hpp"
4*d7be555eSGeorge Liu 
5*d7be555eSGeorge Liu #include <boost/asio/io_context.hpp>
6*d7be555eSGeorge Liu #include <boost/asio/steady_timer.hpp>
7*d7be555eSGeorge Liu 
8*d7be555eSGeorge Liu #include <memory>
9*d7be555eSGeorge Liu #include <stdexcept>
10*d7be555eSGeorge Liu 
11*d7be555eSGeorge Liu class NVMeContext : public std::enable_shared_from_this<NVMeContext>
12*d7be555eSGeorge Liu {
13*d7be555eSGeorge Liu   public:
NVMeContext(boost::asio::io_context & io,int rootBus)14*d7be555eSGeorge Liu     NVMeContext(boost::asio::io_context& io, int rootBus) :
15*d7be555eSGeorge Liu         scanTimer(io), rootBus(rootBus), pollCursor(sensors.end())
16*d7be555eSGeorge Liu     {
17*d7be555eSGeorge Liu         if (rootBus < 0)
18*d7be555eSGeorge Liu         {
19*d7be555eSGeorge Liu             throw std::invalid_argument(
20*d7be555eSGeorge Liu                 "Invalid root bus: Bus ID must not be negative");
21*d7be555eSGeorge Liu         }
22*d7be555eSGeorge Liu     }
23*d7be555eSGeorge Liu 
~NVMeContext()24*d7be555eSGeorge Liu     virtual ~NVMeContext()
25*d7be555eSGeorge Liu     {
26*d7be555eSGeorge Liu         scanTimer.cancel();
27*d7be555eSGeorge Liu     }
28*d7be555eSGeorge Liu 
addSensor(const std::shared_ptr<NVMeSensor> & sensor)29*d7be555eSGeorge Liu     void addSensor(const std::shared_ptr<NVMeSensor>& sensor)
30*d7be555eSGeorge Liu     {
31*d7be555eSGeorge Liu         sensors.emplace_back(sensor);
32*d7be555eSGeorge Liu     }
33*d7be555eSGeorge Liu 
34*d7be555eSGeorge Liu     std::optional<std::shared_ptr<NVMeSensor>>
getSensorAtPath(const std::string & path)35*d7be555eSGeorge Liu         getSensorAtPath(const std::string& path)
36*d7be555eSGeorge Liu     {
37*d7be555eSGeorge Liu         for (auto& sensor : sensors)
38*d7be555eSGeorge Liu         {
39*d7be555eSGeorge Liu             if (sensor->configurationPath == path)
40*d7be555eSGeorge Liu             {
41*d7be555eSGeorge Liu                 return sensor;
42*d7be555eSGeorge Liu             }
43*d7be555eSGeorge Liu         }
44*d7be555eSGeorge Liu 
45*d7be555eSGeorge Liu         return std::nullopt;
46*d7be555eSGeorge Liu     }
47*d7be555eSGeorge Liu 
48*d7be555eSGeorge Liu     // Post-condition: The sensor list does not contain the provided sensor
49*d7be555eSGeorge Liu     // Post-condition: pollCursor is a valid iterator for the sensor list
removeSensor(const std::shared_ptr<NVMeSensor> & sensor)50*d7be555eSGeorge Liu     void removeSensor(const std::shared_ptr<NVMeSensor>& sensor)
51*d7be555eSGeorge Liu     {
52*d7be555eSGeorge Liu         // Locate the sensor that we're removing in the sensor list
53*d7be555eSGeorge Liu         auto found = std::find(sensors.begin(), sensors.end(), sensor);
54*d7be555eSGeorge Liu 
55*d7be555eSGeorge Liu         // If we failed to find the sensor in the list the post-condition is
56*d7be555eSGeorge Liu         // already satisfied
57*d7be555eSGeorge Liu         if (found == sensors.end())
58*d7be555eSGeorge Liu         {
59*d7be555eSGeorge Liu             return;
60*d7be555eSGeorge Liu         }
61*d7be555eSGeorge Liu 
62*d7be555eSGeorge Liu         // We've found the sensor in the list
63*d7be555eSGeorge Liu 
64*d7be555eSGeorge Liu         // If we're not actively polling the sensor list, then remove the sensor
65*d7be555eSGeorge Liu         if (pollCursor == sensors.end())
66*d7be555eSGeorge Liu         {
67*d7be555eSGeorge Liu             sensors.erase(found);
68*d7be555eSGeorge Liu             return;
69*d7be555eSGeorge Liu         }
70*d7be555eSGeorge Liu 
71*d7be555eSGeorge Liu         // We're actively polling the sensor list
72*d7be555eSGeorge Liu 
73*d7be555eSGeorge Liu         // If we're not polling the specific sensor that has been removed, then
74*d7be555eSGeorge Liu         // remove the sensor
75*d7be555eSGeorge Liu         if (*pollCursor != *found)
76*d7be555eSGeorge Liu         {
77*d7be555eSGeorge Liu             sensors.erase(found);
78*d7be555eSGeorge Liu             return;
79*d7be555eSGeorge Liu         }
80*d7be555eSGeorge Liu 
81*d7be555eSGeorge Liu         // We're polling the sensor that is being removed
82*d7be555eSGeorge Liu 
83*d7be555eSGeorge Liu         // Remove the sensor and update the poll cursor so the cursor remains
84*d7be555eSGeorge Liu         // valid
85*d7be555eSGeorge Liu         pollCursor = sensors.erase(found);
86*d7be555eSGeorge Liu     }
87*d7be555eSGeorge Liu 
close()88*d7be555eSGeorge Liu     virtual void close()
89*d7be555eSGeorge Liu     {
90*d7be555eSGeorge Liu         scanTimer.cancel();
91*d7be555eSGeorge Liu     }
92*d7be555eSGeorge Liu 
93*d7be555eSGeorge Liu     virtual void pollNVMeDevices() = 0;
94*d7be555eSGeorge Liu 
95*d7be555eSGeorge Liu     virtual void readAndProcessNVMeSensor() = 0;
96*d7be555eSGeorge Liu 
97*d7be555eSGeorge Liu     virtual void processResponse(std::shared_ptr<NVMeSensor>& sensor, void* msg,
98*d7be555eSGeorge Liu                                  size_t len) = 0;
99*d7be555eSGeorge Liu 
100*d7be555eSGeorge Liu   protected:
101*d7be555eSGeorge Liu     boost::asio::steady_timer scanTimer;
102*d7be555eSGeorge Liu     int rootBus; // Root bus for this drive
103*d7be555eSGeorge Liu     std::list<std::shared_ptr<NVMeSensor>> sensors;
104*d7be555eSGeorge Liu     std::list<std::shared_ptr<NVMeSensor>>::iterator pollCursor;
105*d7be555eSGeorge Liu };
106*d7be555eSGeorge Liu 
107*d7be555eSGeorge Liu using NVMEMap = boost::container::flat_map<int, std::shared_ptr<NVMeContext>>;
108*d7be555eSGeorge Liu 
109*d7be555eSGeorge Liu NVMEMap& getNVMEMap();
110