From dc2b6128fa95d4de1d371e160de2e8dd87a48d3b Mon Sep 17 00:00:00 2001
From: Kowshik Prakasam and Pieter Noordhuis
 <pair+kowshik+pieter@pivotallabs.com>
Date: Mon, 4 Mar 2013 12:57:14 -0800
Subject: [PATCH] Test and fix increasing/decreasing memory limit

---
 .../warden/container/features/mem_limit.rb    | 18 +++++++--
 warden/spec/container/linux_spec.rb           | 39 ++++++++++++++++---
 2 files changed, 48 insertions(+), 9 deletions(-)

diff --git a/warden/lib/warden/container/features/mem_limit.rb b/warden/lib/warden/container/features/mem_limit.rb
index 04e9ae7b..1476b8c6 100644
--- a/warden/lib/warden/container/features/mem_limit.rb
+++ b/warden/lib/warden/container/features/mem_limit.rb
@@ -94,9 +94,21 @@ def limit_memory(limit_in_bytes)
           # when the oom notifier is registered.
           start_oom_notifier_if_needed
 
-          ["memory.limit_in_bytes", "memory.memsw.limit_in_bytes"].each do |path|
-            File.open(File.join(cgroup_path(:memory), path), 'w') do |f|
-              f.write(limit_in_bytes.to_s)
+          # The memory limit may be increased or decreased. The fields that are
+          # set have the following invariant:
+          #
+          #   memory.limit_in_bytes <= memory.memsw.limit_in_bytes
+          #
+          # If the limit is increased and memory.limit_in_bytes is set first,
+          # the invariant may not hold. Similarly, if the limit is decreased
+          # and memory.memsw.limit_in_bytes is set first, the invariant may not
+          # hold. However, one of the two fields will always be set
+          # successfully. To mitigate this, both limits are written twice.
+          2.times do
+            ["memory.limit_in_bytes", "memory.memsw.limit_in_bytes"].each do |path|
+              File.open(File.join(cgroup_path(:memory), path), 'w') do |f|
+                f.write(limit_in_bytes.to_s)
+              end
             end
           end
         end
diff --git a/warden/spec/container/linux_spec.rb b/warden/spec/container/linux_spec.rb
index b55f44b4..e8433515 100644
--- a/warden/spec/container/linux_spec.rb
+++ b/warden/spec/container/linux_spec.rb
@@ -196,13 +196,40 @@ def trigger_oom
       response.limit_in_bytes.should == 9223372036854775807
     end
 
-    it "sets `memory.limit_in_bytes` in the correct cgroup" do
-      hund_mb = 100 * 1024 * 1024
-      response = limit_memory(:limit_in_bytes => hund_mb)
-      response.limit_in_bytes.should == hund_mb
+    describe "setting limits" do
+      def integer_from_memory_cgroup(file)
+        File.read(File.join("/sys/fs/cgroup/memory", "instance-#{@handle}", file)).to_i
+      end
+
+      let(:hundred_mb) { 100 * 1024 * 1024 }
+
+      before do
+        response = limit_memory(:limit_in_bytes => hundred_mb)
+        response.limit_in_bytes.should == hundred_mb
+      end
+
+      it "sets `memory.limit_in_bytes`" do
+        integer_from_memory_cgroup("memory.limit_in_bytes").should == hundred_mb
+      end
+
+      it "sets `memory.memsw.limit_in_bytes`" do
+        integer_from_memory_cgroup("memory.memsw.limit_in_bytes").should == hundred_mb
+      end
 
-      raw_lim = File.read(File.join("/sys/fs/cgroup/memory", "instance-#{@handle}", "memory.limit_in_bytes"))
-      raw_lim.to_i.should == hund_mb
+      describe "increasing limits" do
+        before do
+          response = limit_memory(:limit_in_bytes => 2 * hundred_mb)
+          response.limit_in_bytes.should == 2 * hundred_mb
+        end
+
+        it "sets `memory.limit_in_bytes`" do
+          integer_from_memory_cgroup("memory.limit_in_bytes").should == 2 * hundred_mb
+        end
+
+        it "sets `memory.memsw.limit_in_bytes`" do
+          integer_from_memory_cgroup("memory.memsw.limit_in_bytes").should == 2 * hundred_mb
+        end
+      end
     end
 
     def self.it_should_stop_container_when_an_oom_event_occurs