From 1382bf8720b23b45fc2950f598540dc0a0ba3ece Mon Sep 17 00:00:00 2001
From: Braden MacDonald <mail@bradenm.com>
Date: Wed, 18 Sep 2019 07:27:46 -0700
Subject: [PATCH] Save user state for Blockstore XBlocks in CSM, clean up CSM a
 bit (#21630)

This commit introduces the changes needed for XBlocks in Blockstore to save
their user state into CSM. Before this commit, all student state for Blockstore
blocks was ephemeral (in-process dict store).

Notes:

* The main risk factor of this PR is that it adds non-course keys to the
  course_id field in CSM. If any code (like analytics?) reads course keys
  directly out of CSM and doesn't have graceful handling for key types it
  doesn't recognize, it could cause an issue. With the included changes to
  opaque-keys, calling CourseKey.from_string(...) on these values will raise
  InvalidKeyError since they're not CourseKeys. (But calling
  LearningContextKey.from_string(...) will work for both course and library
  keys.)
* This commit introduces a slight regression for the Studio view of XBlocks in
  Blockstore content libraries: their state is now lost from request to request.
  I have a follow up PR to give them a proper studio-appropriate state store,
  but I want to review it separately so it doesn't hold up this PR and we can
  test this PR on its own.
---
 .../xmodule/video_module/video_module.py      |   3 +-
 .../migrations/0012_adjust_fields.py          |  32 ++++
 lms/djangoapps/courseware/model_data.py       |   9 +-
 lms/djangoapps/courseware/models.py           |  24 +--
 lms/envs/common.py                            |   7 +
 lms/envs/test.py                              |   7 +
 .../core/djangoapps/content_libraries/api.py  |   7 +-
 .../content_libraries/library_context.py      |   8 +-
 .../tests/test_content_libraries.py           |   1 +
 .../tests/test_student_state.py               | 162 ++++++++++++++++++
 openedx/core/djangoapps/xblock/api.py         |   2 +-
 openedx/core/djangoapps/xblock/apps.py        |   6 +-
 .../learning_context/learning_context.py      |   4 +-
 .../xblock/runtime/blockstore_field_data.py   |  42 +++--
 .../core/djangoapps/xblock/runtime/runtime.py | 117 +++++++++----
 15 files changed, 351 insertions(+), 80 deletions(-)
 create mode 100644 lms/djangoapps/courseware/migrations/0012_adjust_fields.py
 create mode 100644 openedx/core/djangoapps/content_libraries/tests/test_student_state.py

diff --git a/common/lib/xmodule/xmodule/video_module/video_module.py b/common/lib/xmodule/xmodule/video_module/video_module.py
index a19fcf2aac2..01d5a770720 100644
--- a/common/lib/xmodule/xmodule/video_module/video_module.py
+++ b/common/lib/xmodule/xmodule/video_module/video_module.py
@@ -23,7 +23,6 @@ from operator import itemgetter
 import six
 from django.conf import settings
 from lxml import etree
-from opaque_keys.edx.keys import CourseKey
 from opaque_keys.edx.locator import AssetLocator
 from web_fragments.fragment import Fragment
 from xblock.completable import XBlockCompletionMode
@@ -201,7 +200,7 @@ class VideoBlock(
         return waffle_flags()[DEPRECATE_YOUTUBE].is_enabled(self.location.course_key)
 
     def youtube_disabled_for_course(self):
-        if not isinstance(self.location.course_key, CourseKey):
+        if not self.location.context_key.is_course:
             return False  # Only courses have this flag
         if CourseYoutubeBlockedFlag.feature_enabled(self.location.course_key):
             return True
diff --git a/lms/djangoapps/courseware/migrations/0012_adjust_fields.py b/lms/djangoapps/courseware/migrations/0012_adjust_fields.py
new file mode 100644
index 00000000000..186cc83021e
--- /dev/null
+++ b/lms/djangoapps/courseware/migrations/0012_adjust_fields.py
@@ -0,0 +1,32 @@
+# -*- coding: utf-8 -*-
+# Generated by Django 1.11.23 on 2019-09-08 04:54
+#
+# This migration does not produce any actual database changes; it only affects
+# the python code. You can confirm this with:
+#    ./manage.py lms sqlmigrate courseware 0012_adjust_fields
+from __future__ import unicode_literals
+
+from django.conf import settings
+from django.db import migrations, models
+import opaque_keys.edx.django.models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+        ('courseware', '0011_csm_id_bigint'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='studentmodule',
+            name='course_id',
+            field=opaque_keys.edx.django.models.LearningContextKeyField(db_index=True, max_length=255),
+        ),
+        migrations.AlterField(
+            model_name='studentmodule',
+            name='module_type',
+            field=models.CharField(db_index=True, max_length=32),
+        ),
+    ]
diff --git a/lms/djangoapps/courseware/model_data.py b/lms/djangoapps/courseware/model_data.py
index 5a3fe4f6baf..7c8c092c676 100644
--- a/lms/djangoapps/courseware/model_data.py
+++ b/lms/djangoapps/courseware/model_data.py
@@ -33,7 +33,7 @@ from contracts import contract, new_contract
 from django.db import DatabaseError, IntegrityError, transaction
 from opaque_keys.edx.asides import AsideUsageKeyV1, AsideUsageKeyV2
 from opaque_keys.edx.block_types import BlockTypeKeyV1
-from opaque_keys.edx.keys import CourseKey
+from opaque_keys.edx.keys import LearningContextKey
 from xblock.core import XBlockAside
 from xblock.exceptions import InvalidScopeError, KeyValueMultiSaveError
 from xblock.fields import Scope, UserScope
@@ -703,7 +703,7 @@ class FieldDataCache(object):
         else:
             self.asides = asides
 
-        assert isinstance(course_id, CourseKey)
+        assert isinstance(course_id, LearningContextKey)
         self.course_id = course_id
         self.user = user
         self.read_only = read_only
@@ -997,13 +997,14 @@ def set_score(user_id, usage_key, score, max_score):
     Set the score and max_score for the specified user and xblock usage.
     """
     created = False
-    kwargs = {"student_id": user_id, "module_state_key": usage_key, "course_id": usage_key.course_key}
+    kwargs = {"student_id": user_id, "module_state_key": usage_key, "course_id": usage_key.context_key}
     try:
         with transaction.atomic():
             student_module, created = StudentModule.objects.get_or_create(
                 defaults={
                     'grade': score,
                     'max_grade': max_score,
+                    'module_type': usage_key.block_type,
                 },
                 **kwargs
             )
@@ -1012,7 +1013,7 @@ def set_score(user_id, usage_key, score, max_score):
         log.exception(
             u'set_score: IntegrityError for student %s - course_id %s - usage_key %s having '
             u'score %d and max_score %d',
-            str(user_id), usage_key.course_key, usage_key, score, max_score
+            str(user_id), usage_key.context_key, usage_key, score, max_score
         )
         student_module = StudentModule.objects.get(**kwargs)
 
diff --git a/lms/djangoapps/courseware/models.py b/lms/djangoapps/courseware/models.py
index 17bccea4888..c25d142bfe7 100644
--- a/lms/djangoapps/courseware/models.py
+++ b/lms/djangoapps/courseware/models.py
@@ -25,7 +25,7 @@ from django.db import models
 from django.db.models.signals import post_save
 from django.utils.translation import ugettext_lazy as _
 from model_utils.models import TimeStampedModel
-from opaque_keys.edx.django.models import BlockTypeKeyField, CourseKeyField, UsageKeyField
+from opaque_keys.edx.django.models import BlockTypeKeyField, CourseKeyField, LearningContextKeyField, UsageKeyField
 from courseware.fields import UnsignedBigIntAutoField
 from six import text_type
 from six.moves import range
@@ -79,33 +79,25 @@ class ChunkingManager(models.Manager):
 
 class StudentModule(models.Model):
     """
-    Keeps student state for a particular module in a particular course.
+    Keeps student state for a particular XBlock usage and particular student.
+
+    Called Module since it was originally used for XModule state.
 
     .. no_pii:
     """
     objects = ChunkingManager()
-    MODEL_TAGS = ['course_id', 'module_type']
-
-    # For a homework problem, contains a JSON
-    # object consisting of state
-    MODULE_TYPES = (('problem', 'problem'),
-                    ('video', 'video'),
-                    ('html', 'html'),
-                    ('course', 'course'),
-                    ('chapter', 'Section'),
-                    ('sequential', 'Subsection'),
-                    ('library_content', 'Library Content'))
 
     id = UnsignedBigIntAutoField(primary_key=True)  # pylint: disable=invalid-name
 
-    ## These three are the key for the object
-    module_type = models.CharField(max_length=32, choices=MODULE_TYPES, default='problem', db_index=True)
+    ## The XBlock/XModule type (e.g. "problem")
+    module_type = models.CharField(max_length=32, db_index=True)
 
     # Key used to share state. This is the XBlock usage_id
     module_state_key = UsageKeyField(max_length=255, db_column='module_id')
     student = models.ForeignKey(User, db_index=True, db_constraint=False, on_delete=models.CASCADE)
 
-    course_id = CourseKeyField(max_length=255, db_index=True)
+    # The learning context of the usage_key (usually a course ID, but may be a library or something else)
+    course_id = LearningContextKeyField(max_length=255, db_index=True)
 
     class Meta(object):
         app_label = "courseware"
diff --git a/lms/envs/common.py b/lms/envs/common.py
index eadc7b842bf..b4c31c98506 100644
--- a/lms/envs/common.py
+++ b/lms/envs/common.py
@@ -499,6 +499,13 @@ DATABASE_ROUTERS = [
 ############################ Cache Configuration ###############################
 
 CACHES = {
+    'blockstore': {
+        'KEY_PREFIX': 'blockstore',
+        'KEY_FUNCTION': 'util.memcache.safe_key',
+        'LOCATION': ['localhost:11211'],
+        'TIMEOUT': '86400',  # This data should be long-lived for performance, BundleCache handles invalidation
+        'BACKEND': 'django.core.cache.backends.memcached.MemcachedCache',
+    },
     'course_structure_cache': {
         'KEY_PREFIX': 'course_structure',
         'KEY_FUNCTION': 'util.memcache.safe_key',
diff --git a/lms/envs/test.py b/lms/envs/test.py
index 2557ebeb101..8375c9baa6c 100644
--- a/lms/envs/test.py
+++ b/lms/envs/test.py
@@ -233,6 +233,13 @@ CACHES = {
     'course_structure_cache': {
         'BACKEND': 'django.core.cache.backends.dummy.DummyCache',
     },
+    # Blockstore caching tests require a cache that actually works:
+    'blockstore': {
+        'KEY_PREFIX': 'blockstore',
+        'KEY_FUNCTION': 'util.memcache.safe_key',
+        'LOCATION': 'edx_loc_mem_cache',
+        'BACKEND': 'django.core.cache.backends.locmem.LocMemCache',
+    },
 }
 
 ############################### BLOCKSTORE #####################################
diff --git a/openedx/core/djangoapps/content_libraries/api.py b/openedx/core/djangoapps/content_libraries/api.py
index be3d1b809b2..20cbdfa32ae 100644
--- a/openedx/core/djangoapps/content_libraries/api.py
+++ b/openedx/core/djangoapps/content_libraries/api.py
@@ -15,7 +15,6 @@ import six
 from xblock.core import XBlock
 from xblock.exceptions import XBlockNotFoundError
 
-from cms.djangoapps.contentstore.views.helpers import xblock_type_display_name
 from openedx.core.djangoapps.content_libraries.library_bundle import LibraryBundle
 from openedx.core.djangoapps.xblock.api import get_block_display_name, load_block
 from openedx.core.djangoapps.xblock.learning_context.manager import get_learning_context_impl
@@ -271,7 +270,7 @@ def get_library_block(usage_key):
     """
     assert isinstance(usage_key, LibraryUsageLocatorV2)
     lib_context = get_learning_context_impl(usage_key)
-    def_key = lib_context.definition_for_usage(usage_key)
+    def_key = lib_context.definition_for_usage(usage_key, force_draft=DRAFT_NAME)
     if def_key is None:
         raise ContentLibraryBlockNotFound(usage_key)
     lib_bundle = LibraryBundle(usage_key.lib_key, def_key.bundle_uuid, draft_name=DRAFT_NAME)
@@ -441,6 +440,10 @@ def get_allowed_block_types(library_key):  # pylint: disable=unused-argument
     library. For now, the result is the same regardless of which library is
     specified, but that may change in the future.
     """
+    # This import breaks in the LMS so keep it here. The LMS doesn't generally
+    # use content libraries APIs directly but some tests may want to use them to
+    # create libraries and then test library learning or course-library integration.
+    from cms.djangoapps.contentstore.views.helpers import xblock_type_display_name
     # TODO: return support status and template options
     # See cms/djangoapps/contentstore/views/component.py
     block_types = sorted(name for name, class_ in XBlock.load_classes())
diff --git a/openedx/core/djangoapps/content_libraries/library_context.py b/openedx/core/djangoapps/content_libraries/library_context.py
index 69130311a66..034ce430109 100644
--- a/openedx/core/djangoapps/content_libraries/library_context.py
+++ b/openedx/core/djangoapps/content_libraries/library_context.py
@@ -55,7 +55,7 @@ class LibraryContextImpl(LearningContext):
         # TODO: implement permissions
         return True
 
-    def definition_for_usage(self, usage_key):
+    def definition_for_usage(self, usage_key, **kwargs):
         """
         Given a usage key for an XBlock in this context, return the
         BundleDefinitionLocator which specifies the actual XBlock definition
@@ -69,7 +69,11 @@ class LibraryContextImpl(LearningContext):
             bundle_uuid = bundle_uuid_for_library_key(library_key)
         except ContentLibrary.DoesNotExist:
             return None
-        bundle = LibraryBundle(library_key, bundle_uuid, self.use_draft)
+        if 'force_draft' in kwargs:
+            use_draft = kwargs['force_draft']
+        else:
+            use_draft = self.use_draft
+        bundle = LibraryBundle(library_key, bundle_uuid, use_draft)
         return bundle.definition_for_usage(usage_key)
 
     def usage_for_child_include(self, parent_usage, parent_definition, parsed_include):
diff --git a/openedx/core/djangoapps/content_libraries/tests/test_content_libraries.py b/openedx/core/djangoapps/content_libraries/tests/test_content_libraries.py
index f87a00cf9fe..7a74702d898 100644
--- a/openedx/core/djangoapps/content_libraries/tests/test_content_libraries.py
+++ b/openedx/core/djangoapps/content_libraries/tests/test_content_libraries.py
@@ -29,6 +29,7 @@ URL_BLOCK_GET_HANDLER_URL = '/api/xblock/v2/xblocks/{block_key}/handler_url/{han
 
 
 @unittest.skipUnless(settings.RUN_BLOCKSTORE_TESTS, "Requires a running Blockstore server")
+@unittest.skipUnless(settings.ROOT_URLCONF == "cms.urls", "Content Libraries REST API is only available in Studio")
 class ContentLibrariesTest(APITestCase):
     """
     Test for Blockstore-based Content Libraries
diff --git a/openedx/core/djangoapps/content_libraries/tests/test_student_state.py b/openedx/core/djangoapps/content_libraries/tests/test_student_state.py
new file mode 100644
index 00000000000..d4bcb121a77
--- /dev/null
+++ b/openedx/core/djangoapps/content_libraries/tests/test_student_state.py
@@ -0,0 +1,162 @@
+# -*- coding: utf-8 -*-
+"""
+Test that the Blockstore-based XBlock runtime can store and retrieve student
+state for XBlocks when learners access blocks directly in a library context,
+if the library allows direct learning.
+"""
+from __future__ import absolute_import, division, print_function, unicode_literals
+import unittest
+
+from django.conf import settings
+from django.test import TestCase
+from organizations.models import Organization
+from xblock.core import XBlock, Scope
+from xblock import fields
+
+from openedx.core.djangoapps.content_libraries import api as library_api
+from openedx.core.djangoapps.xblock import api as xblock_api
+from openedx.core.lib import blockstore_api
+from student.tests.factories import UserFactory
+
+
+class UserStateTestBlock(XBlock):
+    """
+    Block for testing variously scoped XBlock fields.
+    """
+    BLOCK_TYPE = "user-state-test"
+    has_score = False
+
+    display_name = fields.String(scope=Scope.content, name='User State Test Block')
+    # User-specific fields:
+    user_str = fields.String(scope=Scope.user_state, default='default value')  # This usage, one user
+    uss_str = fields.String(scope=Scope.user_state_summary, default='default value')  # This usage, all users
+    pref_str = fields.String(scope=Scope.preferences, default='default value')  # Block type, one user
+    user_info_str = fields.String(scope=Scope.user_info, default='default value')  # All blocks, one user
+
+
+@unittest.skipUnless(settings.RUN_BLOCKSTORE_TESTS, "Requires a running Blockstore server")
+# We can remove the line below to enable this in Studio once we implement a session-backed
+# field data store which we can use for both studio users and anonymous users
+@unittest.skipUnless(settings.ROOT_URLCONF == "lms.urls", "Student State is only saved in the LMS")
+class ContentLibraryXBlockUserStateTest(TestCase):
+    """
+    Test that the Blockstore-based XBlock runtime can store and retrieve student
+    state for XBlocks when learners access blocks directly in a library context,
+    if the library allows direct learning.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        super(ContentLibraryXBlockUserStateTest, cls).setUpClass()
+        # Create a couple students that the tests can use
+        cls.student_a = UserFactory.create(username="Alice", email="alice@example.com")
+        cls.student_b = UserFactory.create(username="Bob", email="bob@example.com")
+        # Create a collection using Blockstore API directly only because there
+        # is not yet any Studio REST API for doing so:
+        cls.collection = blockstore_api.create_collection("Content Library Test Collection")
+        # Create an organization
+        cls.organization = Organization.objects.create(
+            name="Content Libraries Tachyon Exploration & Survey Team",
+            short_name="CL-TEST",
+        )
+        cls.library = library_api.create_library(
+            collection_uuid=cls.collection.uuid,
+            org=cls.organization,
+            slug="state-test-lib",
+            title="Student State Test Lib",
+            description="",
+        )
+
+    @XBlock.register_temp_plugin(UserStateTestBlock, UserStateTestBlock.BLOCK_TYPE)
+    def test_default_values(self):
+        """
+        Test that a user sees the default field values at first
+        """
+        block_metadata = library_api.create_library_block(self.library.key, UserStateTestBlock.BLOCK_TYPE, "b1")
+        block_usage_key = block_metadata.usage_key
+        library_api.publish_changes(self.library.key)
+
+        block_alice = xblock_api.load_block(block_usage_key, self.student_a)
+
+        self.assertEqual(block_alice.scope_ids.user_id, self.student_a.id)
+        self.assertEqual(block_alice.user_str, 'default value')
+        self.assertEqual(block_alice.uss_str, 'default value')
+        self.assertEqual(block_alice.pref_str, 'default value')
+        self.assertEqual(block_alice.user_info_str, 'default value')
+
+    @XBlock.register_temp_plugin(UserStateTestBlock, UserStateTestBlock.BLOCK_TYPE)
+    def test_modify_state_directly(self):
+        """
+        Test that we can modify user-specific XBlock fields directly in Python
+        """
+        # Create two XBlocks, block1 and block2
+        block1_metadata = library_api.create_library_block(self.library.key, UserStateTestBlock.BLOCK_TYPE, "b2-1")
+        block1_usage_key = block1_metadata.usage_key
+        block2_metadata = library_api.create_library_block(self.library.key, UserStateTestBlock.BLOCK_TYPE, "b2-2")
+        block2_usage_key = block2_metadata.usage_key
+        library_api.publish_changes(self.library.key)
+
+        # Alice changes all the fields of block1:
+        block1_alice = xblock_api.load_block(block1_usage_key, self.student_a)
+        block1_alice.user_str = 'Alice was here'
+        block1_alice.uss_str = 'Alice was here (USS)'
+        block1_alice.pref_str = 'Alice was here (prefs)'
+        block1_alice.user_info_str = 'Alice was here (user info)'
+        block1_alice.save()
+
+        # Now load it back and expect the same field data:
+        block1_alice = xblock_api.load_block(block1_usage_key, self.student_a)
+
+        self.assertEqual(block1_alice.scope_ids.user_id, self.student_a.id)
+        self.assertEqual(block1_alice.user_str, 'Alice was here')
+        self.assertEqual(block1_alice.uss_str, 'Alice was here (USS)')
+        self.assertEqual(block1_alice.pref_str, 'Alice was here (prefs)')
+        self.assertEqual(block1_alice.user_info_str, 'Alice was here (user info)')
+
+        # Now load a different block for Alice:
+        block2_alice = xblock_api.load_block(block2_usage_key, self.student_a)
+        # User state should be default:
+        self.assertEqual(block2_alice.user_str, 'default value')
+        # User state summary should be default:
+        self.assertEqual(block2_alice.uss_str, 'default value')
+        # But prefs and user info should be shared:
+        self.assertEqual(block2_alice.pref_str, 'Alice was here (prefs)')
+        self.assertEqual(block2_alice.user_info_str, 'Alice was here (user info)')
+
+        # Now load the first block, block1, for Bob:
+        block1_bob = xblock_api.load_block(block1_usage_key, self.student_b)
+
+        self.assertEqual(block1_bob.scope_ids.user_id, self.student_b.id)
+        self.assertEqual(block1_bob.user_str, 'default value')
+        self.assertEqual(block1_bob.uss_str, 'Alice was here (USS)')
+        self.assertEqual(block1_bob.pref_str, 'default value')
+        self.assertEqual(block1_bob.user_info_str, 'default value')
+
+    @XBlock.register_temp_plugin(UserStateTestBlock, UserStateTestBlock.BLOCK_TYPE)
+    def test_independent_instances(self):
+        """
+        Test that independent instances of the same block don't share field data
+        until .save() and re-loading, even when they're using the same runtime.
+        """
+        block_metadata = library_api.create_library_block(self.library.key, UserStateTestBlock.BLOCK_TYPE, "b3")
+        block_usage_key = block_metadata.usage_key
+        library_api.publish_changes(self.library.key)
+
+        block_instance1 = xblock_api.load_block(block_usage_key, self.student_a)
+        block_instance2 = block_instance1.runtime.get_block(block_usage_key)
+
+        # We could assert that both instances of the block have the same runtime
+        # instance, but that's an implementation detail. The main point of this
+        # test is just to make sure there's never any surprises when reading
+        # field data out of an XBlock, because of other instances of the same
+        # block.
+
+        block_instance1.user_str = 'changed to this'
+        self.assertNotEqual(block_instance1.user_str, block_instance2.user_str)
+
+        block_instance1.save()
+        self.assertNotEqual(block_instance1.user_str, block_instance2.user_str)
+
+        block_instance2 = block_instance1.runtime.get_block(block_usage_key)
+        # Now they should be equal, because we've saved and re-loaded instance2:
+        self.assertEqual(block_instance1.user_str, block_instance2.user_str)
diff --git a/openedx/core/djangoapps/xblock/api.py b/openedx/core/djangoapps/xblock/api.py
index 12890bf0ecf..3f08999fc44 100644
--- a/openedx/core/djangoapps/xblock/api.py
+++ b/openedx/core/djangoapps/xblock/api.py
@@ -79,7 +79,7 @@ def load_block(usage_key, user):
     # set to 3.
     # field_overrides = context_impl.get_field_overrides(usage_key)
 
-    runtime = get_runtime_system().get_runtime(user_id=user.id if user else None)
+    runtime = get_runtime_system().get_runtime(user=user)
 
     return runtime.get_block(usage_key)
 
diff --git a/openedx/core/djangoapps/xblock/apps.py b/openedx/core/djangoapps/xblock/apps.py
index d5d3d1cc061..3490aabae15 100644
--- a/openedx/core/djangoapps/xblock/apps.py
+++ b/openedx/core/djangoapps/xblock/apps.py
@@ -59,8 +59,7 @@ class LmsXBlockAppConfig(XBlockAppConfig):
         editing XBlock content in the LMS
         """
         return dict(
-            authored_data_store=BlockstoreFieldData(),
-            student_data_store=KvsFieldData(kvs=DictKeyValueStore()),
+            student_data_mode='persisted',
         )
 
     def get_site_root_url(self):
@@ -86,8 +85,7 @@ class StudioXBlockAppConfig(XBlockAppConfig):
         editing XBlock content in Studio
         """
         return dict(
-            authored_data_store=BlockstoreFieldData(),
-            student_data_store=KvsFieldData(kvs=DictKeyValueStore()),
+            student_data_mode='ephemeral',
         )
 
     def get_site_root_url(self):
diff --git a/openedx/core/djangoapps/xblock/learning_context/learning_context.py b/openedx/core/djangoapps/xblock/learning_context/learning_context.py
index 457e7f22256..b4e275e13ed 100644
--- a/openedx/core/djangoapps/xblock/learning_context/learning_context.py
+++ b/openedx/core/djangoapps/xblock/learning_context/learning_context.py
@@ -51,7 +51,7 @@ class LearningContext(object):
         """
         return False
 
-    def definition_for_usage(self, usage_key):
+    def definition_for_usage(self, usage_key, **kwargs):
         """
         Given a usage key for an XBlock in this context, return the
         BundleDefinitionLocator which specifies the actual XBlock definition
@@ -59,6 +59,8 @@ class LearningContext(object):
 
         usage_key: the UsageKeyV2 subclass used for this learning context
 
+        kwargs: optional additional parameters unique to the learning context
+
         Must return a BundleDefinitionLocator if the XBlock exists in this
         context, or None otherwise.
         """
diff --git a/openedx/core/djangoapps/xblock/runtime/blockstore_field_data.py b/openedx/core/djangoapps/xblock/runtime/blockstore_field_data.py
index 2cd4061d4ba..64f706420d3 100644
--- a/openedx/core/djangoapps/xblock/runtime/blockstore_field_data.py
+++ b/openedx/core/djangoapps/xblock/runtime/blockstore_field_data.py
@@ -27,12 +27,33 @@ MAX_DEFINITIONS_LOADED = 100  # How many of the most recently used XBlocks' fiel
 class BlockInstanceUniqueKey(object):
     """
     An empty object used as a unique key for each XBlock instance, see
-    BlockstoreFieldData._get_active_block(). Every XBlock instance will get a
-    unique one of these keys, even if they are otherwise identical. Its purpose
-    is similar to `id(block)`.
+    get_weak_key_for_block() and BlockstoreFieldData._get_active_block(). Every
+    XBlock instance will get a unique one of these keys, even if they are
+    otherwise identical. Its purpose is similar to `id(block)`.
     """
 
 
+def get_weak_key_for_block(block):
+    """
+    Given an XBlock instance, return an object with the same lifetime as the
+    block, suitable as a key to hold block-specific data in a WeakKeyDictionary.
+    """
+    # We would like to make the XBlock instance 'block' itself the key of
+    # BlockstoreFieldData.active_blocks, so that we have exactly one entry per
+    # XBlock instance in memory, and they'll each be automatically freed by the
+    # WeakKeyDictionary as needed. But because XModules implement
+    # __eq__() in a way that reads all field values, just attempting to use
+    # the block as a dict key here will trigger infinite recursion. So
+    # instead we key the dict on an arbitrary object,
+    # block key = BlockInstanceUniqueKey() which we create here. That way
+    # the weak reference will still cause the entry in the WeakKeyDictionary to
+    # be freed automatically when the block is no longer needed, and we
+    # still get one entry per XBlock instance.
+    if not hasattr(block, '_field_data_key_obj'):
+        block._field_data_key_obj = BlockInstanceUniqueKey()  # pylint: disable=protected-access
+    return block._field_data_key_obj  # pylint: disable=protected-access
+
+
 def get_olx_hash_for_definition_key(def_key):
     """
     Given a BundleDefinitionLocator, which identifies a specific version of an
@@ -123,20 +144,7 @@ class BlockstoreFieldData(FieldData):
         Get the ActiveBlock entry for the specified block, creating it if
         necessary.
         """
-        # We would like to make the XBlock instance 'block' itself the key of
-        # self.active_blocks, so that we have exactly one entry per XBlock
-        # instance in memory, and they'll each be automatically freed by the
-        # WeakKeyDictionary as needed. But because XModules implement
-        # __eq__() in a way that reads all field values, just attempting to use
-        # the block as a dict key here will trigger infinite recursion. So
-        # instead we key the dict on an arbitrary object,
-        # block key = BlockInstanceUniqueKey() which we create here. That way
-        # the weak reference will still cause the entry in self.active_blocks to
-        # be freed automatically when the block is no longer needed, and we
-        # still get one entry per XBlock instance.
-        if not hasattr(block, '_field_data_key_obj'):
-            block._field_data_key_obj = BlockInstanceUniqueKey()  # pylint: disable=protected-access
-        key = block._field_data_key_obj  # pylint: disable=protected-access
+        key = get_weak_key_for_block(block)
         if key not in self.active_blocks:
             self.active_blocks[key] = ActiveBlock(
                 olx_hash=get_olx_hash_for_definition_key(block.scope_ids.def_id),
diff --git a/openedx/core/djangoapps/xblock/runtime/runtime.py b/openedx/core/djangoapps/xblock/runtime/runtime.py
index 77476a1ad8c..face80ce4c0 100644
--- a/openedx/core/djangoapps/xblock/runtime/runtime.py
+++ b/openedx/core/djangoapps/xblock/runtime/runtime.py
@@ -4,15 +4,18 @@ Common base classes for all new XBlock runtimes.
 from __future__ import absolute_import, division, print_function, unicode_literals
 import logging
 
+from django.contrib.auth import get_user_model
 from django.utils.lru_cache import lru_cache
 from six.moves.urllib.parse import urljoin  # pylint: disable=import-error
 from xblock.exceptions import NoSuchServiceError
 from xblock.field_data import SplitFieldData
 from xblock.fields import Scope
-from xblock.runtime import Runtime, NullI18nService, MemoryIdManager
+from xblock.runtime import DictKeyValueStore, KvsFieldData, NullI18nService, MemoryIdManager, Runtime
 from web_fragments.fragment import Fragment
 
+from courseware.model_data import DjangoKeyValueStore, FieldDataCache
 from openedx.core.djangoapps.xblock.apps import get_xblock_app_config
+from openedx.core.djangoapps.xblock.runtime.blockstore_field_data import BlockstoreFieldData
 from openedx.core.lib.xblock_utils import xblock_local_resource_url
 from xmodule.errortracker import make_error_tracker
 from .id_managers import OpaqueKeyReader
@@ -20,6 +23,7 @@ from .shims import RuntimeShim, XBlockShim
 
 
 log = logging.getLogger(__name__)
+User = get_user_model()
 
 
 class XBlockRuntime(RuntimeShim, Runtime):
@@ -37,8 +41,7 @@ class XBlockRuntime(RuntimeShim, Runtime):
     # ** Do not add any XModule compatibility code to this class **
     # Add it to RuntimeShim instead, to help keep legacy code isolated.
 
-    def __init__(self, system, user_id):
-        # type: (XBlockRuntimeSystem, int) -> None
+    def __init__(self, system, user):
         super(XBlockRuntime, self).__init__(
             id_reader=system.id_reader,
             mixins=(
@@ -52,7 +55,10 @@ class XBlockRuntime(RuntimeShim, Runtime):
             id_generator=system.id_generator,
         )
         self.system = system
-        self.user_id = user_id
+        self.user = user
+        self.user_id = user.id if self.user else None  # Must be set as a separate attribute since base class sets it
+        self.block_field_datas = {}  # dict of FieldData stores for our loaded XBlocks. Key is the block's scope_ids.
+        self.django_field_data_caches = {}  # dict of FieldDataCache objects for XBlock with database-based user state
 
     def handler_url(self, block, handler_name, suffix='', query='', thirdparty=False):
         """
@@ -114,12 +120,67 @@ class XBlockRuntime(RuntimeShim, Runtime):
         declaration = block.service_declaration(service_name)
         if declaration is None:
             raise NoSuchServiceError("Service {!r} was not requested.".format(service_name))
-        # Special case handling for some services:
-        service = self.system.get_service(block.scope_ids, service_name)
+        # Most common service is field-data so check that first:
+        if service_name == "field-data":
+            if block.scope_ids not in self.block_field_datas:
+                try:
+                    self.block_field_datas[block.scope_ids] = self._init_field_data_for_block(block)
+                except:
+                    # Don't try again pointlessly every time another field is accessed
+                    self.block_field_datas[block.scope_ids] = None
+                    raise
+            return self.block_field_datas[block.scope_ids]
+        # Check if the XBlockRuntimeSystem wants to handle this:
+        service = self.system.get_service(block, service_name)
+        # Otherwise, fall back to the base implementation which loads services
+        # defined in the constructor:
         if service is None:
             service = super(XBlockRuntime, self).service(block, service_name)
         return service
 
+    def _init_field_data_for_block(self, block):
+        """
+        Initialize the FieldData implementation for the specified XBlock
+        """
+        if self.user is None:
+            # No user is specified, so we want to throw an error if anything attempts to read/write user-specific fields
+            student_data_store = None
+        elif self.user.is_anonymous:
+            # The user is anonymous. Future work will support saving their state
+            # in a cache or the django session but for now just use a highly
+            # ephemeral dict.
+            student_data_store = KvsFieldData(kvs=DictKeyValueStore())
+        elif self.system.student_data_mode == XBlockRuntimeSystem.STUDENT_DATA_EPHEMERAL:
+            # We're in an environment like Studio where we want to let the
+            # author test blocks out but not permanently save their state.
+            # This in-memory dict will typically only persist for one
+            # request-response cycle, so we need to soon replace it with a store
+            # that puts the state into a cache or the django session.
+            student_data_store = KvsFieldData(kvs=DictKeyValueStore())
+        else:
+            # Use database-backed field data (i.e. store user_state in StudentModule)
+            context_key = block.scope_ids.usage_id.context_key
+            if context_key not in self.django_field_data_caches:
+                field_data_cache = FieldDataCache(
+                    [block], course_id=context_key, user=self.user, asides=None, read_only=False,
+                )
+                self.django_field_data_caches[context_key] = field_data_cache
+            else:
+                field_data_cache = self.django_field_data_caches[context_key]
+                field_data_cache.add_descriptors_to_cache([block])
+            student_data_store = KvsFieldData(kvs=DjangoKeyValueStore(field_data_cache))
+
+        return SplitFieldData({
+            Scope.content: self.system.authored_data_store,
+            Scope.settings: self.system.authored_data_store,
+            Scope.parent: self.system.authored_data_store,
+            Scope.children: self.system.authored_data_store,
+            Scope.user_state_summary: student_data_store,
+            Scope.user_state: student_data_store,
+            Scope.user_info: student_data_store,
+            Scope.preferences: student_data_store,
+        })
+
     def render(self, block, view_name, context=None):
         """
         Render a specific view of an XBlock.
@@ -157,11 +218,13 @@ class XBlockRuntimeSystem(object):
     """
     ANONYMOUS_USER = 'anon'  # Special value passed to handler_url() methods
 
+    STUDENT_DATA_EPHEMERAL = 'ephemeral'
+    STUDENT_DATA_PERSISTED = 'persisted'
+
     def __init__(
         self,
         handler_url,  # type: (Callable[[UsageKey, str, Union[int, ANONYMOUS_USER]], str]
-        authored_data_store,  # type: FieldData
-        student_data_store,  # type: FieldData
+        student_data_mode,  # type: Union[STUDENT_DATA_EPHEMERAL, STUDENT_DATA_PERSISTED]
         runtime_class,  # type: XBlockRuntime
     ):
         """
@@ -175,34 +238,28 @@ class XBlockRuntimeSystem(object):
                 )
                 If user_id is ANONYMOUS_USER, the handler should execute without
                 any user-scoped fields.
-            authored_data_store: A FieldData instance used to retrieve/write
-                any fields with UserScope.NONE
-            student_data_store: A FieldData instance used to retrieve/write
-                any fields with UserScope.ONE or UserScope.ALL
+            student_data_mode: Specifies whether student data should be kept
+                in a temporary in-memory store (e.g. Studio) or persisted
+                forever in the database.
+            runtime_class: What runtime to use, e.g. BlockstoreXBlockRuntime
         """
         self.handler_url = handler_url
         self.id_reader = OpaqueKeyReader()
         self.id_generator = MemoryIdManager()  # We don't really use id_generator until we need to support asides
         self.runtime_class = runtime_class
-        self.authored_data_store = authored_data_store
-        self.field_data = SplitFieldData({
-            Scope.content: authored_data_store,
-            Scope.settings: authored_data_store,
-            Scope.parent: authored_data_store,
-            Scope.children: authored_data_store,
-            Scope.user_state_summary: student_data_store,
-            Scope.user_state: student_data_store,
-            Scope.user_info: student_data_store,
-            Scope.preferences: student_data_store,
-        })
-
+        self.authored_data_store = BlockstoreFieldData()
+        assert student_data_mode in (self.STUDENT_DATA_EPHEMERAL, self.STUDENT_DATA_PERSISTED)
+        self.student_data_mode = student_data_mode
         self._error_trackers = {}
 
-    def get_runtime(self, user_id):
-        # type: (int) -> XBlockRuntime
-        return self.runtime_class(self, user_id)
+    def get_runtime(self, user):
+        """
+        Get the XBlock runtime for the specified Django user. The user can be
+        a regular user, an AnonymousUser, or None.
+        """
+        return self.runtime_class(self, user)
 
-    def get_service(self, scope_ids, service_name):
+    def get_service(self, block, service_name):
         """
         Get a runtime service
 
@@ -210,10 +267,8 @@ class XBlockRuntimeSystem(object):
         or if this method returns None, they may come from the
         XBlockRuntime.
         """
-        if service_name == "field-data":
-            return self.field_data
         if service_name == 'error_tracker':
-            return self.get_error_tracker_for_context(scope_ids.usage_id.context_key)
+            return self.get_error_tracker_for_context(block.scope_ids.usage_id.context_key)
         return None  # None means see if XBlockRuntime offers this service
 
     @lru_cache(maxsize=32)
-- 
GitLab