Class: Kcl::Checkpointer

Inherits:
Object
  • Object
show all
Defined in:
lib/kcl/checkpointer.rb

Constant Summary collapse

DYNAMO_DB_LEASE_PRIMARY_KEY = 
'shard_id'.freeze
DYNAMO_DB_LEASE_OWNER_KEY = 
'assigned_to'.freeze
DYNAMO_DB_LEASE_TIMEOUT_KEY = 
'lease_timeout'.freeze
DYNAMO_DB_CHECKPOINT_SEQUENCE_NUMBER_KEY = 
'checkpoint'.freeze
DYNAMO_DB_PARENT_SHARD_KEY = 
'parent_shard_id'.freeze

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(config) ⇒ Checkpointer

Returns a new instance of Checkpointer.

Parameters:



13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
 
# File 'lib/kcl/checkpointer.rb', line 13

def initialize(config)
  @dynamodb = Kcl::Proxies::DynamoDbProxy.new(config)
  @table_name = config.dynamodb_table_name

  return if @dynamodb.exists?(@table_name)
  @dynamodb.create_table(
    @table_name,
    [{
      attribute_name: DYNAMO_DB_LEASE_PRIMARY_KEY,
      attribute_type: 'S'
    }],
    [{
      attribute_name: DYNAMO_DB_LEASE_PRIMARY_KEY,
      key_type: 'HASH'
    }],
    {
      read_capacity_units: config.dynamodb_read_capacity,
      write_capacity_units: config.dynamodb_write_capacity
    }
  )
  Kcl.logger.info("Created DynamoDB table: #{@table_name}")
end

Instance Attribute Details

#dynamodb ⇒ Object (readonly)

Returns the value of attribute dynamodb.



10
11
12
 
# File 'lib/kcl/checkpointer.rb', line 10

def dynamodb
  @dynamodb
end

Instance Method Details

#fetch_checkpoint(shard) ⇒ Kcl::Workers::ShardInfo

Retrieves the checkpoint for the given shard

Returns:



39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
 
# File 'lib/kcl/checkpointer.rb', line 39

def fetch_checkpoint(shard)
  checkpoint = @dynamodb.get_item(
    @table_name,
    { "#{DYNAMO_DB_LEASE_PRIMARY_KEY}" => shard.shard_id }
  )
  return shard if checkpoint.nil?

  if checkpoint[DYNAMO_DB_CHECKPOINT_SEQUENCE_NUMBER_KEY]
    shard.checkpoint = checkpoint[DYNAMO_DB_CHECKPOINT_SEQUENCE_NUMBER_KEY]
  end
  if checkpoint[DYNAMO_DB_LEASE_OWNER_KEY]
    shard.assigned_to = checkpoint[DYNAMO_DB_LEASE_OWNER_KEY]
  end
  Kcl.logger.info("Retrieves checkpoint of shard at #{shard.to_h}")

  shard
end

#lease(shard, next_assigned_to) ⇒ Kcl::Workers::ShardInfo

Attempt to gain a lock on the given shard

Returns:



85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
 
# File 'lib/kcl/checkpointer.rb', line 85

def lease(shard, next_assigned_to)
  now = Time.now.utc
  next_lease_timeout = now + Kcl.config.dynamodb_failover_seconds

  checkpoint = @dynamodb.get_item(
    @table_name,
    { "#{DYNAMO_DB_LEASE_PRIMARY_KEY}" => shard.shard_id }
  )
  assigned_to   = checkpoint && checkpoint[DYNAMO_DB_LEASE_OWNER_KEY]
  lease_timeout = checkpoint && checkpoint[DYNAMO_DB_LEASE_TIMEOUT_KEY]

  if assigned_to && lease_timeout
    if now > Time.parse(lease_timeout) && assigned_to != next_assigned_to
      raise Kcl::Errors::LeaseNotAquiredError
    end
    condition_expression = 'shard_id = :shard_id AND assigned_to = :assigned_to AND lease_timeout = :lease_timeout'
    expression_attributes = {
      ':shard_id' => shard.shard_id,
      ':assigned_to' => assigned_to,
      ':lease_timeout' => lease_timeout
    }
    Kcl.logger.info("Attempting to get a lock for shard: #{shard.to_h}")
  else
    condition_expression = 'attribute_not_exists(assigned_to)'
    expression_attributes = nil
  end

  item = {
    "#{DYNAMO_DB_LEASE_PRIMARY_KEY}" => shard.shard_id,
    "#{DYNAMO_DB_LEASE_OWNER_KEY}"   => next_assigned_to,
    "#{DYNAMO_DB_LEASE_TIMEOUT_KEY}" => next_lease_timeout.to_s
  }
  if shard.checkpoint != ''
    item[DYNAMO_DB_CHECKPOINT_SEQUENCE_NUMBER_KEY] = shard.checkpoint
  end
  if shard.parent_shard_id > 0
    item[DYNAMO_DB_PARENT_SHARD_KEY] = shard.parent_shard_id
  end

  result = @dynamodb.conditional_update_item(
    @table_name,
    item,
    condition_expression,
    expression_attributes
  )
  if result
    shard.assigned_to   = next_assigned_to
    shard.lease_timeout = next_lease_timeout
    Kcl.logger.info("Get lease for shard at #{shard.to_h}")
  else
    Kcl.logger.info("Failed to get lease for shard at #{shard.to_h}")
  end

  shard
end

#remove_lease(shard) ⇒ Kcl::Workers::ShardInfo

Remove the shard entry

Returns:



144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
 
# File 'lib/kcl/checkpointer.rb', line 144

def remove_lease(shard)
  result = @dynamodb.remove_item(
    @table_name,
    { "#{DYNAMO_DB_LEASE_PRIMARY_KEY}" => shard.shard_id }
  )
  if result
    shard.assigned_to   = nil
    shard.checkpoint    = nil
    shard.lease_timeout = nil
    Kcl.logger.info("Remove lease for shard at #{shard.to_h}")
  else
    Kcl.logger.info("Failed to remove lease for shard at #{shard.to_h}")
  end

  shard
end

#remove_lease_owner(shard) ⇒ Kcl::Workers::ShardInfo

Remove lease owner for the shard entry

Returns:



164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
 
# File 'lib/kcl/checkpointer.rb', line 164

def remove_lease_owner(shard)
  result = @dynamodb.update_item(
    @table_name,
    { "#{DYNAMO_DB_LEASE_PRIMARY_KEY}" => shard.shard_id },
    "remove #{DYNAMO_DB_LEASE_OWNER_KEY}"
  )
  if result
    shard.assigned_to = nil
    Kcl.logger.info("Remove lease owner for shard at #{shard.to_h}")
  else
    Kcl.logger.info("Failed to remove lease owner for shard at #{shard.to_h}")
  end

  shard
end

#update_checkpoint(shard) ⇒ Kcl::Workers::ShardInfo

Write the checkpoint for the given shard

Returns:



60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
 
# File 'lib/kcl/checkpointer.rb', line 60

def update_checkpoint(shard)
  item = {
    "#{DYNAMO_DB_LEASE_PRIMARY_KEY}" => shard.shard_id,
    "#{DYNAMO_DB_CHECKPOINT_SEQUENCE_NUMBER_KEY}" => shard.checkpoint,
    "#{DYNAMO_DB_LEASE_OWNER_KEY}" => shard.assigned_to,
    "#{DYNAMO_DB_LEASE_TIMEOUT_KEY}" => shard.lease_timeout.to_s
  }
  if shard.parent_shard_id > 0
    item[DYNAMO_DB_PARENT_SHARD_KEY] = shard.parent_shard_id
  end

  result = @dynamodb.put_item(@table_name, item)
  if result
    Kcl.logger.info("Write checkpoint of shard at #{shard.to_h}")
  else
    Kcl.logger.info("Failed to write checkpoint for shard at #{shard.to_h}")
  end

  shard
end