Add shard plugin (#877)

* Add shard plugin

This adds a shard_seed that people can then shard on ala:

```
def in_shard(node, shard, shard_size)
  node['shard_seed'] % shard_size <= shard
end
```

Which then enables:

```
if in_shard(node, 4, 100)
  # some new stuff
else
  # some old stuff
end
```

Which is roughly how we shard things. Other code to utilize this to follow.

* Add 'hostname' vs 'fqdn' options; make function non-global

* require library, allow multiple sources

* Add mac support, change defaults, fix config

* Fix style

* Add tests, fix DMI lookup

* More style fixes

* Add failure in the event of unrecognized sources

* More style fixes

* Refector code for less repetition. Thanks @coderanger

* Style fixes

2 files changed, 142 insertions, 0 deletions

diff --git a/lib/ohai/plugins/shard.rb b/lib/ohai/plugins/shard.rb
new file mode 100644
index 00000000..7cc14148
--- /dev/null
+++ b/lib/ohai/plugins/shard.rb
@@ -0,0 +1,82 @@
+#
+# Author:: Phil Dibowitz <phil@ipom.com>
+# Copyright:: Copyright (c) 2016 Facebook, Inc.
+# License:: Apache License, Version 2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+require "digest/md5"
+
+Ohai.plugin(:ShardSeed) do
+  depends "hostname", "dmi", "machine_id", "machinename"
+  provides "shard_seed"
+
+  def get_dmi_property(dmi, thing)
+    %w{system base_board chassis}.each do |section|
+      unless dmi[section][thing].strip.empty?
+        return dmi[section][thing]
+      end
+    end
+  end
+
+  def default_sources
+    [:machinename, :serial, :uuid]
+  end
+
+  # Common sources go here. Put sources that need to be different per-platform
+  # under their collect_data block.
+  def create_seed(&block)
+    sources = Ohai.config[:plugin][:shard_seed][:sources] || default_sources
+    data = ""
+    sources.each do |src|
+      data << case src
+              when :fqdn
+                fqdn
+              when :hostname
+                hostname
+              when :machine_id
+                machine_id
+              when :machinename
+                machinename
+              else
+                yield(src)
+              end
+    end
+    shard_seed Digest::MD5.hexdigest(data)[0...7].to_i(16)
+  end
+
+  collect_data(:darwin) do
+    create_seed do |src|
+      case src
+      when :serial
+        hardware["serial_number"]
+      when :uuid
+        hardware["platform_UUID"]
+      end
+    end
+  end
+
+  collect_data(:linux) do
+    create_seed do |src|
+      case src
+      when :serial
+        get_dmi_property(dmi, :serial_number)
+      when :uuid
+        get_dmi_property(dmi, :uuid)
+      else
+        raise "No such shard_seed source: #{src}"
+      end
+    end
+  end
+end
diff --git a/spec/unit/plugins/shard_spec.rb b/spec/unit/plugins/shard_spec.rb
new file mode 100644
index 00000000..5d6ef066
--- /dev/null
+++ b/spec/unit/plugins/shard_spec.rb
@@ -0,0 +1,60 @@
+#
+# Author:: Davide Cavalca <dcavalca@fb.com>
+# Copyright:: Copyright (c) 2016 Facebook
+# License:: Apache License, Version 2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+require "digest/md5"
+require File.expand_path(File.dirname(__FILE__) + "/../../spec_helper.rb")
+
+describe Ohai::System, "shard plugin" do
+  let(:plugin) { get_plugin("shard") }
+  let(:fqdn) { "somehost004.someregion.somecompany.com" }
+  let(:uuid) { "48555CF4-5BB1-21D9-BC4C-E8B73DDE5801" }
+  let(:serial) { "234du3m4i498xdjr2" }
+  let(:machine_id) { "0a1f869f457a4c8080ab19faf80af9cc" }
+  let(:machinename) { "somehost004" }
+
+  before(:each) do
+    allow(plugin).to receive(:collect_os).and_return(:linux)
+    plugin["machinename"] = machinename
+    plugin["machine_id"] = machine_id
+    plugin["fqdn"] = fqdn
+    plugin["dmi"] = { "system" => {} }
+    plugin["dmi"]["system"]["uuid"] = uuid
+    plugin["dmi"]["system"]["serial_number"] = serial
+    allow(plugin).to receive(:collect_os).and_return(:linux)
+  end
+
+  it "should provide a shard with a default-safe set of sources" do
+    plugin.run
+    result = Digest::MD5.hexdigest(
+      "#{machinename}#{serial}#{uuid}"
+    )[0...7].to_i(16)
+    expect(plugin[:shard_seed]).to eq(result)
+  end
+
+  it "should provide a shard with a configured source" do
+    Ohai.config[:plugin][:shard_seed][:sources] = [:fqdn]
+    plugin.run
+    result = Digest::MD5.hexdigest(fqdn)[0...7].to_i(16)
+    expect(plugin[:shard_seed]).to eq(result)
+  end
+
+  it "fails on an unrecognized source" do
+    Ohai.config[:plugin][:shard_seed][:sources] = [:GreatGooglyMoogly]
+    expect { plugin.run }.to raise_error(RuntimeError)
+  end
+end