From 5e9c76d6dfdcd81f0be11f1f7009083746154d76 Mon Sep 17 00:00:00 2001
From: Holger Just <hello@holgerjust.de>
Date: Wed, 13 Jun 2018 22:33:45 +0200
Subject: [PATCH] Allow to rotate log files using a date-based pattern

We now support two different modes of file rotation at the same time:

* auto_reopen can be used to automatically reopen a logfile at the
  original location if the file was moved or deleted from the filesystem
* rotate can be used to write to a rotate file which can be reopened /
  created based on Date pattern.

The user can now decide whether they want to use an external logrotate
command or use internal rotation with Rackstash instead.
---
 lib/rackstash/adapter/file.rb       | 112 ++++++++++++++++---
 spec/rackstash/adapter/file_spec.rb | 163 +++++++++++++++++++++++++---
 2 files changed, 240 insertions(+), 35 deletions(-)

diff --git a/lib/rackstash/adapter/file.rb b/lib/rackstash/adapter/file.rb
index b30157e..2fb798a 100644
--- a/lib/rackstash/adapter/file.rb
+++ b/lib/rackstash/adapter/file.rb
@@ -64,9 +64,23 @@ module Rackstash
     class File < BaseAdapter
       register_for ::String, ::Pathname, 'file'
 
-      # @return [String] the absolute path to the log file
+      # @return [String] the absolute path to the currently opened log file
       attr_reader :path
 
+      # @return [String] the absolute path to the originally defined log file.
+      #   Depending on the {#rotate} setting, the final log file might have a
+      #   date-based suffix added before its file extension. Use {#path} to
+      #   get the full path of the currently opened log file.
+      attr_reader :base_path
+
+      # @return [String, Proc, nil] date pattern for the file suffix used for
+      #   auto-rotated log files. The pattern is used with `Date#strftime` to
+      #   determine the file suffix for the current rotate file. When setting a
+      #   `Proc`, it is expected to return the currently final log file suffix
+      #   (not just a date pattern). When setting the value to `nil`, the log
+      #   file is not rotated.
+      attr_reader :rotate
+
       def self.from_uri(uri)
         uri = URI(uri)
 
@@ -91,17 +105,19 @@ module Rackstash
       # [`::File.expand_path`](https://ruby-doc.org/core/File.html#method-c-expand_path)
       # for details.
       #
-      # @param path [String, Pathname] the path to the logfile
-      # @param auto_reopen [Boolean] set to `true` to automatically reopen the
-      #   log file (and potentially create a new one) if we detect that the
-      #   current log file was moved or deleted, e.g. due to an external
-      #   logrotate run
-      def initialize(path, auto_reopen: true)
-        @path = ::File.expand_path(path).freeze
-        @auto_reopen = !!auto_reopen
+      # @param path [String, Pathname] the path to the logfile. Depending on the
+      #   `rotate` setting, the final log file might have a date-based suffix
+      #   added before its file extension.
+      # @param auto_reopen (see #auto_reopen=)
+      # @param rotate (see #rotate=)
+      def initialize(path, auto_reopen: true, rotate: nil)
+        @base_path = ::File.expand_path(path).freeze
+
+        self.auto_reopen = auto_reopen
+        self.rotate = rotate
 
         @mutex = Mutex.new
-        open_file
+        open_file(rotated_path)
       end
 
       # @return [Boolean] if `true`, the logfile will be automatically reopened
@@ -110,6 +126,40 @@ module Rackstash
         @auto_reopen
       end
 
+      # @param auto_reopen [Boolean] set to `true` to automatically reopen the
+      #   log file (and potentially create a new one) if we detect that the
+      #   current log file was moved or deleted, e.g. due to an external
+      #   logrotate run
+      def auto_reopen=(auto_reopen)
+        @auto_reopen = !!auto_reopen
+      end
+
+      # @param rotate [String, Proc, nil] date pattern for the file suffix used
+      #   for auto-rotated log files. When giving a `String` here, it is
+      #   interpreted as a pattern for the `Date#strftime` method. In addition
+      #   to that, we accept the following names: `"daily"`, `"weekly"`, and
+      #   `"monthly"` for pre-defined suffixes. When giving a `Proc`, it is
+      #   expected to return the final suffix on call (i.e. not just a
+      #   `Date#strftime` pattern but the actual file suffix). When defining a
+      #   rotate pattern, each log event is written to a file with the resulting
+      #   suffix added before its file extension.
+      def rotate=(rotate)
+        @rotate = case rotate
+        when :daily, 'daily'.freeze
+          '%Y-%m-%d'.freeze
+        when :weekly, 'weekly'.freeze
+          '%G-w%V'.freeze
+        when :monthly, 'monthly'.freeze
+          '%Y-%m'.freeze
+        when String
+          rotate.dup.freeze
+        when Proc, nil
+          rotate
+        else
+          raise ArgumentError, "Invalid rotate specification: #{rotate.inspect}"
+        end
+      end
+
       # Write a single log line with a trailing newline character to the open
       # file. If {#auto_reopen?} is `true`, we will reopen the file object
       # before the write if we detect that the file was moved, e.g., from an
@@ -128,7 +178,7 @@ module Rackstash
         return if line.empty?
 
         @mutex.synchronize do
-          auto_reopen
+          rotate_file
           @file.syswrite(line)
         end
         nil
@@ -155,7 +205,7 @@ module Rackstash
       # @return [nil]
       def reopen
         @mutex.synchronize do
-          reopen_file
+          reopen_file rotated_path
         end
         nil
       end
@@ -164,16 +214,17 @@ module Rackstash
 
       # Reopen the log file if the original {#path} does not reference the
       # opened file anymore (e.g. because it was moved or deleted)
-      def auto_reopen
+      def auto_reopen!
         return unless @auto_reopen
+        return unless @file && @path
 
         return if @file.closed?
         return if ::File.identical?(@file, @path)
 
-        reopen_file
+        reopen_file(@path)
       end
 
-      def open_file
+      def open_file(path)
         dirname = ::File.dirname(path)
         FileUtils.mkdir_p(dirname) unless ::File.exist?(dirname)
 
@@ -184,13 +235,40 @@ module Rackstash
         file.binmode
         file.sync = true
 
+        @path = path
         @file = file
         nil
       end
 
-      def reopen_file
+      def reopen_file(path)
         @file.close rescue nil
-        open_file
+        open_file(path)
+      end
+
+      def rotate_file
+        path = rotated_path
+
+        if path == @path
+          auto_reopen!
+        else
+          reopen_file(path)
+        end
+      end
+
+      def rotated_path
+        suffix = case @rotate
+        when String
+          Date.today.strftime(@rotate)
+        when Proc
+          @rotate.call.to_s
+        else
+          EMPTY_STRING
+        end
+
+        return @base_path if suffix.empty?
+
+        suffix = ".#{suffix}"
+        @base_path.sub(/\A(.*?)(\.[^.\/]+)?\z/) { "#{$1}#{suffix}#{$2}" }
       end
     end
   end
diff --git a/spec/rackstash/adapter/file_spec.rb b/spec/rackstash/adapter/file_spec.rb
index 5ad4efe..e40bc3e 100644
--- a/spec/rackstash/adapter/file_spec.rb
+++ b/spec/rackstash/adapter/file_spec.rb
@@ -18,28 +18,32 @@ RSpec.describe Rackstash::Adapter::File do
   let(:adapter) { described_class.new(logfile.path, **adapter_args) }
 
   after(:each) do
+    # Cleanup
+    FileUtils.rm_f Dir.glob("#{logfile.path}.*")
     logfile.close
     logfile.unlink
   end
 
   describe 'from_uri' do
     it 'creates a File adapter instance' do
-      expect(described_class.from_uri('file:/tmp/file_spec.log'))
+      expect(described_class.from_uri("file:#{logfile.path}"))
         .to be_instance_of described_class
-      expect(described_class.from_uri('file:///tmp/file_spec.log'))
+      expect(described_class.from_uri("file://#{logfile.path}"))
         .to be_instance_of described_class
     end
 
-    it 'sets the path from the URI path' do
-      expect(described_class.from_uri('file:/tmp/file_spec.log').path)
-        .to eql '/tmp/file_spec.log'
-      expect(described_class.from_uri('file:///tmp/file_spec.log').path)
-        .to eql '/tmp/file_spec.log'
+    it 'sets the base_path from the URI path' do
+      expect(described_class.from_uri("file:#{logfile.path}").base_path)
+        .to eql logfile.path
+      expect(described_class.from_uri("file://#{logfile.path}").base_path)
+        .to eql logfile.path
     end
 
     it 'sets optional attributes' do
-      expect(described_class.from_uri('file:/tmp/file_spec.log?auto_reopen=false').auto_reopen?)
-        .to eql false
+      adapter = described_class.from_uri('file:/tmp/file_spec.log?rotate=monthly&auto_reopen=false')
+
+      expect(adapter.rotate).to eql '%Y-%m'
+      expect(adapter.auto_reopen?).to eql false
     end
 
     it 'only accepts file URIs' do
@@ -53,13 +57,13 @@ RSpec.describe Rackstash::Adapter::File do
 
   describe '#initialize' do
     it 'accepts a String' do
-      expect(described_class.new(logfile.path).path)
+      expect(described_class.new(logfile.path).base_path)
         .to eql(logfile.path)
         .and be_a String
     end
 
     it 'accepts a Pathname' do
-      expect(described_class.new(Pathname.new(logfile.path)).path)
+      expect(described_class.new(Pathname.new(logfile.path)).base_path)
         .to eql(logfile.path)
         .and be_a String
     end
@@ -76,11 +80,18 @@ RSpec.describe Rackstash::Adapter::File do
 
         adapter = described_class.new File.join(base, 'dir', 'sub', 'test.log')
 
-        expect(adapter.path).to eql File.join(base, 'dir', 'sub', 'test.log')
+        expect(adapter.base_path).to eql File.join(base, 'dir', 'sub', 'test.log')
         expect(File.directory?(File.join(base, 'dir'))).to be true
         expect(File.file?(File.join(base, 'dir', 'sub', 'test.log'))).to be true
       end
     end
+
+    it 'rejects invalid rotate specifications' do
+      expect { described_class.new(logfile.path, rotate: :invalid) }.to raise_error ArgumentError
+      expect { described_class.new(logfile.path, rotate: 42) }.to raise_error ArgumentError
+      expect { described_class.new(logfile.path, rotate: false) }.to raise_error ArgumentError
+      expect { described_class.new(logfile.path, rotate: true) }.to raise_error ArgumentError
+    end
   end
 
   describe '.default_encoder' do
@@ -132,14 +143,14 @@ RSpec.describe Rackstash::Adapter::File do
       let(:adapter_args) { { auto_reopen: true } }
 
       it 'reopens the file if moved' do
-        expect(adapter.auto_reopen?).to be true
+        expect(adapter.auto_reopen?).to eql true
 
         adapter.write('line1')
-        File.rename(logfile.path, "#{logfile.path}.orig")
+        File.rename(logfile.path, "#{logfile.path}.moved")
 
         adapter.write('line2')
 
-        expect(File.read("#{logfile.path}.orig")).to eql "line1\n"
+        expect(File.read("#{logfile.path}.moved")).to eql "line1\n"
         expect(File.read(logfile.path)).to eql "line2\n"
       end
     end
@@ -148,17 +159,133 @@ RSpec.describe Rackstash::Adapter::File do
       let(:adapter_args) { { auto_reopen: false } }
 
       it 'does not reopen the logfile automatically' do
-        expect(adapter.auto_reopen?).to be false
+        expect(adapter.auto_reopen?).to eql false
 
         adapter.write('line1')
-        File.rename(logfile.path, "#{logfile.path}.orig")
+        File.rename(logfile.path, "#{logfile.path}.moved")
 
         adapter.write('line2')
 
-        expect(File.read("#{logfile.path}.orig")).to eql "line1\nline2\n"
+        expect(File.read("#{logfile.path}.moved")).to eql "line1\nline2\n"
         expect(File.exist?(logfile.path)).to be false
       end
     end
+
+    context 'with rotate: :daily' do
+      before do
+        adapter_args[:rotate] = :daily
+      end
+
+      it 'rotates daily' do
+        date1 = Date.new(2017, 11, 13)
+        allow(Date).to receive(:today).and_return(date1)
+
+        adapter.write('line1')
+        expect(adapter.path).to eql "#{logfile.path}.2017-11-13"
+
+        date2 = Date.new(2018, 1, 13)
+        allow(Date).to receive(:today).and_return(date2)
+
+        adapter.write('line2')
+        expect(adapter.path).to eql "#{logfile.path}.2018-01-13"
+
+        expect(File.read "#{logfile.path}.2017-11-13").to eql "line1\n"
+        expect(File.read "#{logfile.path}.2018-01-13").to eql "line2\n"
+      end
+    end
+
+    context 'with rotate: :weekly' do
+      before do
+        adapter_args[:rotate] = :weekly
+      end
+
+      it 'rotates weekly' do
+        date1 = Date.new(2018, 12, 24)
+        allow(Date).to receive(:today).and_return(date1)
+
+        adapter.write('line1')
+        expect(adapter.path).to eql "#{logfile.path}.2018-w52"
+
+        date2 = Date.new(2018, 12, 31)
+        allow(Date).to receive(:today).and_return(date2)
+
+        adapter.write('line2')
+        expect(adapter.path).to eql "#{logfile.path}.2019-w01"
+
+        expect(File.read "#{logfile.path}.2018-w52").to eql "line1\n"
+        expect(File.read "#{logfile.path}.2019-w01").to eql "line2\n"
+      end
+    end
+
+    context 'with rotate: :monthly' do
+      before do
+        adapter_args[:rotate] = :monthly
+      end
+
+      it 'rotates monthly' do
+        date1 = Date.new(2017, 11, 13)
+        allow(Date).to receive(:today).and_return(date1)
+
+        adapter.write('line1')
+        expect(adapter.path).to eql "#{logfile.path}.2017-11"
+
+        date2 = Date.new(2018, 1, 13)
+        allow(Date).to receive(:today).and_return(date2)
+
+        adapter.write('line2')
+        expect(adapter.path).to eql "#{logfile.path}.2018-01"
+
+        expect(File.read "#{logfile.path}.2017-11").to eql "line1\n"
+        expect(File.read "#{logfile.path}.2018-01").to eql "line2\n"
+      end
+    end
+
+    context 'with rotate: PATTERN' do
+      it 'rotates with current year' do
+        adapter_args[:rotate] = 'year-%Y'
+
+        adapter.write('line1')
+        expect(adapter.path).to eql "#{logfile.path}.year-#{Date.today.year}"
+        expect(File.read "#{logfile.path}.year-#{Date.today.year}").to eql "line1\n"
+      end
+
+      it 'rotates with a fixed string' do
+        adapter_args[:rotate] = 'ext'
+
+        adapter.write('line1')
+        expect(adapter.path).to eql "#{logfile.path}.ext"
+
+        adapter.write('line2')
+        expect(adapter.path).to eql "#{logfile.path}.ext"
+
+        expect(File.read "#{logfile.path}.ext").to eql "line1\nline2\n"
+      end
+    end
+
+    context 'with rotate: block' do
+      let(:counter) {
+        Struct.new(:count) do
+          def inc
+            self.count += 1
+          end
+        end.new(0)
+      }
+
+      it 'rotates' do
+        adapter_args[:rotate] = -> { "count_#{counter.inc}" }
+        expect(adapter.path).to eql "#{logfile.path}.count_1"
+
+        adapter.write('line1')
+        expect(adapter.path).to eql "#{logfile.path}.count_2"
+
+        adapter.write('line2')
+        expect(adapter.path).to eql "#{logfile.path}.count_3"
+
+        expect(File.read "#{logfile.path}.count_1").to be_empty
+        expect(File.read "#{logfile.path}.count_2").to eql "line1\n"
+        expect(File.read "#{logfile.path}.count_3").to eql "line2\n"
+      end
+    end
   end
 
   context 'with concurrent processes' do