automatically generate ID for given ipaddress #18

Gemfile

@@ -11,8 +11,8 @@ gem 'rspec-system-puppet',     :require => false
 gem 'serverspec',              :require => false
 gem 'rspec-system-serverspec', :require => false
 # coverage reports will be in release 2.0
-gem 'rspec-puppet', '>= 2.2'
-gem 'rspec', '~> 2.13'
+gem 'rspec-puppet'
+gem 'rspec'
 gem 'metadata-json-lint', :require => false
 
 group :development do

lib/puppet/parser/functions/zookeeper_genid.rb

@@ -0,0 +1,42 @@
+require 'ipaddr'
+
+#
+# ZooKeeper unique ID generator - generates ID for given IP address
+#
+
+module Puppet::Parser::Functions
+  newfunction(:zookeeper_genid, :type => :rvalue, :doc => <<-EOS
+This function generates ZooKeeper ID (1-255) for given IP address
+
+        EOS
+  ) do |args|
+
+    # Arguments: [ipaddress], [strategy]
+
+    if !args[0].nil? && args[0].class != String
+      raise(Puppet::ParseError, "zookeeper_genid() expects an ipaddress, you've given: '" + args[0].class.to_s + "'")
+    end
+
+    # use default ipaddress fact, if no ip given
+    if args[0].nil? || args[0].empty?
+      ipaddress = lookupvar('ipaddress')
+    else
+      ipaddress = IPAddr.new(args[0])
+    end
+    strategy = args[1] if args.size > 1
+    strategy ||= 'mod'
+
+    id = 1
+    case strategy
+    when 'last_digit'
+      ip = ipaddress.to_s
+      idx = ip.rindex('.')
+      return ip[idx+1..-1].to_i
+    when 'mod'
+      # make sure we stay in range 1-255, e.g.: '192.168.1.149' -> 0
+      return (ipaddress.to_i % 255) + 1
+    else
+      raise(Puppet::ParseError, "zookeeper_genid() unknown strategy " + strategy)
+    end
+  end
+end

manifests/init.pp

@@ -13,7 +13,7 @@
 #   class { 'zookeeper': }
 #
 class zookeeper(
-  $id                      = '1',
+  $id                      = undef, # unique ID in ZooKeeper cluster (1-255)
   $datastore               = '/var/lib/zookeeper',
   # datalogstore used to put transaction logs in separate location than snapshots
   $datalogstore            = undef,
@@ -67,6 +67,7 @@
   # donate the matching directives in the [Unit] section
   $systemd_unit_want       = undef,
   $systemd_unit_after      = 'network.target',
+  $id_generator            = 'mod',
 ) inherits ::zookeeper::params {
 
   validate_array($packages)
@@ -93,6 +94,13 @@
     $_manage_service_file = $manage_service_file
   }
 
+  if ($id) {
+    $_id = $id
+  } else {
+    # automatically generate ID
+    $_id = zookeeper_genid($client_ip, $id_generator)
+  }
+
   anchor { 'zookeeper::start': }->
   class { 'zookeeper::install':
     ensure            => $ensure,
@@ -111,7 +119,7 @@
     java_package      => $java_package,
   }->
   class { 'zookeeper::config':
-    id                      => $id,
+    id                      => $_id,
     datastore               => $datastore,
     datalogstore            => $datalogstore,
     initialize_datastore    => $initialize_datastore,

spec/functions/zookeeper_genid_spec.rb

@@ -0,0 +1,42 @@
+#! /usr/bin/env ruby -S rspec
+require 'spec_helper'
+require 'rspec'
+require 'rspec-puppet'
+
+describe 'zookeeper_genid' do
+
+  describe 'convert ipaddress to unique ID in range 1-255' do
+    it 'works with default strategy' do
+      ip = '192.168.1.1'
+      subject.should run.with_params(ip).and_return(108)
+    end
+
+    it 'return valid ID' do
+      ip = '10.0.0.'
+      (1..10).to_a.each do |i|
+        subject.should run.with_params("#{ip}#{i}").and_return(11+i)
+      end
+    end
+
+    it 'should raise an error if run with extra arguments' do
+      subject.should run.with_params(1, 2, 3, 4).and_raise_error(Puppet::ParseError)
+    end
+
+    it 'should raise an error if gets invalid ip address' do
+      subject.should run.with_params('10.0.0.256').and_raise_error(IPAddr::InvalidAddressError)
+    end
+  end
+
+  describe 'last digit strategy' do
+    strategy = 'last_digit'
+    it 'return valid id' do
+      ip = '192.168.1.1'
+      subject.should run.with_params(ip, strategy).and_return(1)
+    end
+
+    it 'works with multiple digits' do
+      ip = '192.168.1.126'
+      subject.should run.with_params(ip, strategy).and_return(126)
+    end
+  end
+end