From d31631ec088c74cf7f190ed2639658394086e2c9 Mon Sep 17 00:00:00 2001
From: Fujimoto Seiji <fujimoto@ceptord.net>
Date: Fri, 1 Jul 2022 14:26:55 +0900
Subject: [PATCH] Add testcase for worker lock functions

Signed-off-by: Fujimoto Seiji <fujimoto@ceptord.net>
---
 lib/fluent/plugin/base.rb | 10 ++++++----
 test/plugin/test_base.rb  | 22 ++++++++++++++++++++++
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/lib/fluent/plugin/base.rb b/lib/fluent/plugin/base.rb
index 28485dcbd5..185d7ec5e5 100644
--- a/lib/fluent/plugin/base.rb
+++ b/lib/fluent/plugin/base.rb
@@ -75,14 +75,16 @@ def multi_workers_ready?
         true
       end
 
+      def worker_lockfile(name)
+        name = name.gsub(/[^a-zA-Z0-9]/, "_")
+        File.join(fluentd_lockdir, "fluentd-#{name}.lock")
+      end
+
       def acquire_worker_lock(name)
         if fluentd_lockdir.nil?
           raise InvalidLockDirectory, "can't acquire lock because FLUENTD_LOCKDIR isn't set"
         end
-
-        name = name.gsub(/[^a-zA-Z0-9]/, "_")
-        lockfile = "fluentd-#{name}.lock"
-        File.open(File.join(fluentd_lockdir, lockfile), "w") do |f|
+        File.open(worker_lockfile(name), "w") do |f|
           f.flock(File::LOCK_EX)
           yield
         end
diff --git a/test/plugin/test_base.rb b/test/plugin/test_base.rb
index 51d55c1b69..307466172b 100644
--- a/test/plugin/test_base.rb
+++ b/test/plugin/test_base.rb
@@ -1,4 +1,5 @@
 require_relative '../helper'
+require 'tmpdir'
 require 'fluent/plugin/base'
 
 module FluentPluginBaseTest
@@ -112,4 +113,25 @@ class FluentPluginBaseTest::DummyPlugin2 < Fluent::Plugin::TestBase
     assert_equal 1, logger.logs.size
     assert{ logger.logs.first.include?("invalid byte sequence is replaced in ") }
   end
+
+  test 'acquire worker lock' do
+    Dir.mktmpdir("test-fluentd-lock-") do |lockdir|
+      ENV['FLUENTD_LOCKDIR'] = lockdir
+      lockfile = @p.worker_lockfile("worker_test")
+
+      @p.acquire_worker_lock("worker_test") do
+        # With LOCK_NB set, flock() returns `false` when the
+        # file is already locked.
+        File.open(lockfile, "w") do |f|
+          assert_equal false, f.flock(File::LOCK_EX|File::LOCK_NB)
+        end
+      end
+
+      # Lock should be release by now. In that case, flock
+      # must return 0.
+      File.open(lockfile, "w") do |f|
+        assert_equal 0, f.flock(File::LOCK_EX|File::LOCK_NB)
+      end
+    end
+  end
 end