diff --git a/per_user_block.module b/per_user_block.module
index 0ae58bf..ce3327c 100644
--- a/per_user_block.module
+++ b/per_user_block.module
@@ -29,8 +29,8 @@ define('PER_USER_BLOCK_CUSTOM_DISABLED', 2);
 function per_user_block_form_user_form_alter(&$form, FormStateInterface $form_state) {
 
   // Load all the block info for the current default theme.
-  $theme = $default_theme = Drupal::config('system.theme')->get('default');
-  $entities = _block_rehash($theme);
+  $theme = $default_theme = \Drupal::config('system.theme')->get('default');
+  $entities = \Drupal::entityManager()->getStorage('block')->loadByProperties(['theme' => $theme]);
 
   // Load any existing user data.
   $account = $form_state->getFormObject()->getEntity();
@@ -79,12 +79,3 @@ function per_user_block_user_form_submit(&$form, FormStateInterface $form_state)
     Drupal::service('user.data')->set('per_user_block', $account->id(), 'blocks', $form_state->getValue('blocks'));
   }
 }
-
-/**
- * Implements hook_entity_info_alter().
- */
-function per_user_block_entity_type_alter(array &$entity_types) {
-  /** @var \Drupal\Core\Config\Entity\ConfigEntityType[] $entity_types */
-  $entity_types['block']
-    ->setClass('Drupal\per_user_block\Entity\Block');
-}
diff --git a/per_user_block.services.yml b/per_user_block.services.yml
index 66c2bda..746eb15 100644
--- a/per_user_block.services.yml
+++ b/per_user_block.services.yml
@@ -1,5 +1,5 @@
 services:
   per_user_block.block_context:
-    class: Drupal\per_user_block\EventSubscriber\BlockContext
+    class: Drupal\per_user_block\ContextProvider\BlockContext
     tags:
-      - { name: 'event_subscriber' }
+      - { name: 'context_provider' }
diff --git a/src/ContextProvider/BlockContext.php b/src/ContextProvider/BlockContext.php
new file mode 100644
index 0000000..3b83346
--- /dev/null
+++ b/src/ContextProvider/BlockContext.php
@@ -0,0 +1,41 @@
+<?php
+
+/**
+ * @file
+ * Contains \Drupal\per_user_block\ContextProvider\BlockContext
+ */
+
+namespace Drupal\per_user_block\ContextProvider;
+
+use Drupal\Core\Cache\CacheableMetadata;
+use Drupal\Core\Plugin\Context\Context;
+use Drupal\Core\Plugin\Context\ContextDefinition;
+use Drupal\Core\Plugin\Context\ContextProviderInterface;
+
+class BlockContext implements ContextProviderInterface {
+
+  /**
+   * {@inheritdoc}
+   */
+  public function getRuntimeContexts(array $unqualified_context_ids) {
+
+    $block = NULL; // @TODO, get this?
+
+    $context = new Context(new ContextDefinition('entity:block'));
+    $context->setContextValue($block);
+
+    $cacheability = new CacheableMetadata();
+    $cacheability->setCacheContexts(['block']);
+    $context->addCacheableDependency($cacheability);
+
+    return ['per_user_block.block' => $context];
+  }
+
+  /**
+   * {@inheritdoc}
+   */
+  public function getAvailableContexts() {
+    return $this->getRuntimeContexts([]);
+  }
+
+}
diff --git a/src/Entity/Block.php b/src/Entity/Block.php
deleted file mode 100644
index 14eb868..0000000
--- a/src/Entity/Block.php
+++ /dev/null
@@ -1,28 +0,0 @@
-<?php
-
-/**
- * @file
- * Contains \Drupal\per_user_block\Entity\Block.
- */
-
-namespace Drupal\per_user_block\Entity;
-
-use Drupal\block\Entity\Block as CoreBlock;
-use Drupal\Core\Plugin\Context\Context;
-use Drupal\Core\Plugin\Context\ContextDefinition;
-
-class Block extends CoreBlock {
-
-  /**
-   * {@inheritdoc}
-   */
-  public function setContexts(array $contexts) {
-    // Store the name in the contexts array using a key that is used in our
-    // BlockContext subscriber.
-    $contexts['per_user_block.block'] = new Context(new ContextDefinition('entity:block'));
-    $contexts['per_user_block.block']->setContextValue($this);
-    $this->contexts = $contexts;
-    return $this;
-  }
-
-}
diff --git a/src/EventSubscriber/BlockContext.php b/src/EventSubscriber/BlockContext.php
deleted file mode 100644
index ddc5fac..0000000
--- a/src/EventSubscriber/BlockContext.php
+++ /dev/null
@@ -1,34 +0,0 @@
-<?php
-
-/**
- * @file
- * Contains \Drupal\per_user_block\EventSubscriber\BlockContext
- */
-
-namespace Drupal\per_user_block\EventSubscriber;
-
-use Drupal\block\Event\BlockContextEvent;
-use Drupal\block\Event\BlockEvents;
-use Drupal\Core\Plugin\Context\Context;
-use Drupal\Core\Plugin\Context\ContextDefinition;
-use Symfony\Component\EventDispatcher\EventSubscriberInterface;
-
-class BlockContext implements EventSubscriberInterface {
-
-  /**
-   * {@inheritdoc}
-   */
-  public function onBlockAdministrativeContext(BlockContextEvent $event) {
-    $context = new Context(new ContextDefinition('entity:block'));
-    $event->setContext('per_user_block.block', $context);
-  }
-
-  /**
-   * {@inheritdoc}
-   */
-  public static function getSubscribedEvents() {
-    $events[BlockEvents::ADMINISTRATIVE_CONTEXT][] = 'onBlockAdministrativeContext';
-    return $events;
-  }
-
-}