From 01d858066a5d037d73b9dc52afa5e5a7971ca9cb Mon Sep 17 00:00:00 2001
From: Frederick Cheung <frederick.cheung@gmail.com>
Date: Sat, 6 Jul 2019 21:41:15 +0100
Subject: [PATCH] Implement faster version for Mac OS

---
 get_process_mem.gemspec       |  1 +
 lib/get_process_mem.rb        | 19 +++++++++++++
 lib/get_process_mem/darwin.rb | 53 +++++++++++++++++++++++++++++++++++
 3 files changed, 73 insertions(+)
 create mode 100644 lib/get_process_mem/darwin.rb

diff --git a/get_process_mem.gemspec b/get_process_mem.gemspec
index a53c403..109e0b0 100644
--- a/get_process_mem.gemspec
+++ b/get_process_mem.gemspec
@@ -19,6 +19,7 @@ Gem::Specification.new do |gem|
   gem.require_paths = ["lib"]
 
   gem.add_development_dependency "sys-proctable", "~> 1.0"
+  gem.add_development_dependency "ffi", "~> 1.0"
   gem.add_development_dependency "rake",  "~> 10.1"
   gem.add_development_dependency "test-unit", "~> 3.1.0"
 end
diff --git a/lib/get_process_mem.rb b/lib/get_process_mem.rb
index 4ef9cfe..ad352eb 100644
--- a/lib/get_process_mem.rb
+++ b/lib/get_process_mem.rb
@@ -31,6 +31,20 @@ class GetProcessMem
     include Sys
   end
 
+  RUNS_ON_DARWIN =  Gem.platforms.detect do |p|
+                      p.is_a?(Gem::Platform) && p.os == 'darwin'
+                    end
+
+  if RUNS_ON_DARWIN
+    begin
+      require 'get_process_mem/darwin'
+    rescue LoadError => e
+      message = "Please add `ffi` to your Gemfile for darwin (macos) machines\n"
+      message << e.message
+      raise e, message
+    end
+  end
+
   def initialize(pid = Process.pid)
     @status_file  = Pathname.new "/proc/#{pid}/status"
     @process_file = Pathname.new "/proc/#{pid}/smaps"
@@ -44,6 +58,7 @@ def linux?
 
   def bytes
     memory =   linux_status_memory if linux?
+    memory ||= darwin_memory if RUNS_ON_DARWIN
     memory ||= ps_memory
   end
 
@@ -102,4 +117,8 @@ def ps_memory
       KB_TO_BYTE * number_to_bigdecimal(mem == "" ? 0 : mem)
     end
   end
+
+  def darwin_memory
+    Darwin.resident_size
+  end
 end
diff --git a/lib/get_process_mem/darwin.rb b/lib/get_process_mem/darwin.rb
new file mode 100644
index 0000000..2fd98cb
--- /dev/null
+++ b/lib/get_process_mem/darwin.rb
@@ -0,0 +1,53 @@
+require 'ffi'
+
+class GetProcessMem
+  class Darwin
+    extend FFI::Library
+    ffi_lib 'c'
+    attach_function :mach_task_self, [], :__darwin_mach_port_t
+    attach_function :task_info,
+                    [
+                      :__darwin_mach_port_t,
+                      :int,     # return selector
+                      :pointer, #pointer to task info
+                      :pointer, #pointer to int (size of structure / bytes filled out)
+                    ],
+                    :int
+
+    class IntPtr < FFI::Struct
+      layout :value, :int
+    end
+
+    class TaskInfo < FFI::Struct
+      layout  :suspend_count, :int32,
+              :virtual_size, :uint64,
+              :resident_size, :uint64,
+              :user_time, :uint64,
+              :system_time, :uint64,
+              :policy, :int32
+    end
+
+    MACH_TASK_BASIC_INFO = 20
+    MACH_TASK_BASIC_INFO_COUNT = TaskInfo.size / FFI.type_size(:uint)
+
+    class << self
+      def resident_size
+        mach_task_info[:resident_size]
+      end
+
+      private
+
+      def mach_task_info
+        data = TaskInfo.new
+        out_count = IntPtr.new
+        out_count[:value] = MACH_TASK_BASIC_INFO_COUNT
+        result = task_info(mach_task_self, MACH_TASK_BASIC_INFO, data, out_count)
+        if result == 0
+          data
+        else
+          raise "task_info returned #{result}"
+        end
+      end
+    end
+  end
+end