WriteRequestBatch.php 9.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267
  1. <?php
  2. namespace Aws\DynamoDb;
  3. use Aws\AwsClientInterface;
  4. use Aws\CommandInterface;
  5. use Aws\CommandPool;
  6. use Aws\Exception\AwsException;
  7. use Aws\ResultInterface;
  8. /**
  9. * The WriteRequestBatch is an object that is capable of efficiently sending
  10. * DynamoDB BatchWriteItem requests from queued up put and delete item requests.
  11. * requests. The batch attempts to send the requests with the fewest requests
  12. * to DynamoDB as possible and also re-queues any unprocessed items to ensure
  13. * that all items are sent.
  14. */
  15. class WriteRequestBatch
  16. {
  17. /** @var DynamoDbClient DynamoDB client used to perform write operations. */
  18. private $client;
  19. /** @var array Configuration options for the batch. */
  20. private $config;
  21. /** @var array Queue of pending put/delete requests in the batch. */
  22. private $queue;
  23. /**
  24. * Creates a WriteRequestBatch object that is capable of efficiently sending
  25. * DynamoDB BatchWriteItem requests from queued up Put and Delete requests.
  26. *
  27. * @param DynamoDbClient $client DynamoDB client used to send batches.
  28. * @param array $config Batch configuration options.
  29. * - table: (string) DynamoDB table used by the batch, this can be
  30. * overridden for each individual put() or delete() call.
  31. * - batch_size: (int) The size of each batch (default: 25). The batch
  32. * size must be between 2 and 25. If you are sending batches of large
  33. * items, you may consider lowering the batch size, otherwise, you
  34. * should use 25.
  35. * - pool_size: (int) This number dictates how many BatchWriteItem
  36. * requests you would like to do in parallel. For example, if the
  37. * "batch_size" is 25, and "pool_size" is 3, then you would send 3
  38. * BatchWriteItem requests at a time, each with 25 items. Please keep
  39. * your throughput in mind when choosing the "pool_size" option.
  40. * - autoflush: (bool) This option allows the batch to automatically
  41. * flush once there are enough items (i.e., "batch_size" * "pool_size")
  42. * in the queue. This defaults to true, so you must set this to false
  43. * to stop autoflush.
  44. * - before: (callable) Executed before every BatchWriteItem operation.
  45. * It should accept an \Aws\CommandInterface object as its argument.
  46. * - error: Executed if an error was encountered executing a,
  47. * BatchWriteItem operation, otherwise errors are ignored. It should
  48. * accept an \Aws\Exception\AwsException as its argument.
  49. *
  50. * @throws \InvalidArgumentException if the batch size is not between 2 and 25.
  51. */
  52. public function __construct(DynamoDbClient $client, array $config = [])
  53. {
  54. // Apply defaults
  55. $config += [
  56. 'table' => null,
  57. 'batch_size' => 25,
  58. 'pool_size' => 1,
  59. 'autoflush' => true,
  60. 'before' => null,
  61. 'error' => null
  62. ];
  63. // Ensure the batch size is valid
  64. if ($config['batch_size'] > 25 || $config['batch_size'] < 2) {
  65. throw new \InvalidArgumentException('"batch_size" must be between 2 and 25.');
  66. }
  67. // Ensure the callbacks are valid
  68. if ($config['before'] && !is_callable($config['before'])) {
  69. throw new \InvalidArgumentException('"before" must be callable.');
  70. }
  71. if ($config['error'] && !is_callable($config['error'])) {
  72. throw new \InvalidArgumentException('"error" must be callable.');
  73. }
  74. // If autoflush is enabled, set the threshold
  75. if ($config['autoflush']) {
  76. $config['threshold'] = $config['batch_size'] * $config['pool_size'];
  77. }
  78. $this->client = $client;
  79. $this->config = $config;
  80. $this->queue = [];
  81. }
  82. /**
  83. * Adds a put item request to the batch.
  84. *
  85. * @param array $item Data for an item to put. Format:
  86. * [
  87. * 'attribute1' => ['type' => 'value'],
  88. * 'attribute2' => ['type' => 'value'],
  89. * ...
  90. * ]
  91. * @param string|null $table The name of the table. This must be specified
  92. * unless the "table" option was provided in the
  93. * config of the WriteRequestBatch.
  94. *
  95. * @return $this
  96. */
  97. public function put(array $item, $table = null)
  98. {
  99. $this->queue[] = [
  100. 'table' => $this->determineTable($table),
  101. 'data' => ['PutRequest' => ['Item' => $item]],
  102. ];
  103. $this->autoFlush();
  104. return $this;
  105. }
  106. /**
  107. * Adds a delete item request to the batch.
  108. *
  109. * @param array $key Key of an item to delete. Format:
  110. * [
  111. * 'key1' => ['type' => 'value'],
  112. * ...
  113. * ]
  114. * @param string|null $table The name of the table. This must be specified
  115. * unless the "table" option was provided in the
  116. * config of the WriteRequestBatch.
  117. *
  118. * @return $this
  119. */
  120. public function delete(array $key, $table = null)
  121. {
  122. $this->queue[] = [
  123. 'table' => $this->determineTable($table),
  124. 'data' => ['DeleteRequest' => ['Key' => $key]],
  125. ];
  126. $this->autoFlush();
  127. return $this;
  128. }
  129. /**
  130. * Flushes the batch by combining all the queued put and delete requests
  131. * into BatchWriteItem commands and executing them. Unprocessed items are
  132. * automatically re-queued.
  133. *
  134. * @param bool $untilEmpty If true, flushing will continue until the queue
  135. * is completely empty. This will make sure that
  136. * unprocessed items are all eventually sent.
  137. *
  138. * @return $this
  139. */
  140. public function flush($untilEmpty = true)
  141. {
  142. // Send BatchWriteItem requests until the queue is empty
  143. $keepFlushing = true;
  144. while ($this->queue && $keepFlushing) {
  145. $commands = $this->prepareCommands();
  146. $pool = new CommandPool($this->client, $commands, [
  147. 'before' => $this->config['before'],
  148. 'concurrency' => $this->config['pool_size'],
  149. 'fulfilled' => function (ResultInterface $result) {
  150. // Re-queue any unprocessed items
  151. if ($result->hasKey('UnprocessedItems')) {
  152. $this->retryUnprocessed($result['UnprocessedItems']);
  153. }
  154. },
  155. 'rejected' => function ($reason) {
  156. if ($reason instanceof AwsException) {
  157. $code = $reason->getAwsErrorCode();
  158. if ($code === 'ProvisionedThroughputExceededException') {
  159. $this->retryUnprocessed($reason->getCommand()['RequestItems']);
  160. } elseif (is_callable($this->config['error'])) {
  161. $this->config['error']($reason);
  162. }
  163. }
  164. }
  165. ]);
  166. $pool->promise()->wait();
  167. $keepFlushing = (bool) $untilEmpty;
  168. }
  169. return $this;
  170. }
  171. /**
  172. * Creates BatchWriteItem commands from the items in the queue.
  173. *
  174. * @return CommandInterface[]
  175. */
  176. private function prepareCommands()
  177. {
  178. // Chunk the queue into batches
  179. $batches = array_chunk($this->queue, $this->config['batch_size']);
  180. $this->queue = [];
  181. // Create BatchWriteItem commands for each batch
  182. $commands = [];
  183. foreach ($batches as $batch) {
  184. $requests = [];
  185. foreach ($batch as $item) {
  186. if (!isset($requests[$item['table']])) {
  187. $requests[$item['table']] = [];
  188. }
  189. $requests[$item['table']][] = $item['data'];
  190. }
  191. $commands[] = $this->client->getCommand(
  192. 'BatchWriteItem',
  193. ['RequestItems' => $requests]
  194. );
  195. }
  196. return $commands;
  197. }
  198. /**
  199. * Re-queues unprocessed results with the correct data.
  200. *
  201. * @param array $unprocessed Unprocessed items from a result.
  202. */
  203. private function retryUnprocessed(array $unprocessed)
  204. {
  205. foreach ($unprocessed as $table => $requests) {
  206. foreach ($requests as $request) {
  207. $this->queue[] = [
  208. 'table' => $table,
  209. 'data' => $request,
  210. ];
  211. }
  212. }
  213. }
  214. /**
  215. * If autoflush is enabled and the threshold is met, flush the batch
  216. */
  217. private function autoFlush()
  218. {
  219. if ($this->config['autoflush']
  220. && count($this->queue) >= $this->config['threshold']
  221. ) {
  222. // Flush only once. Unprocessed items are handled in a later flush.
  223. $this->flush(false);
  224. }
  225. }
  226. /**
  227. * Determine the table name by looking at what was provided and what the
  228. * WriteRequestBatch was originally configured with.
  229. *
  230. * @param string|null $table The table name.
  231. *
  232. * @return string
  233. * @throws \RuntimeException if there was no table specified.
  234. */
  235. private function determineTable($table)
  236. {
  237. $table = $table ?: $this->config['table'];
  238. if (!$table) {
  239. throw new \RuntimeException('There was no table specified.');
  240. }
  241. return $table;
  242. }
  243. }