From 46babbe9e11384f4bde77434e9951d597e1a3c90 Mon Sep 17 00:00:00 2001
From: Anders Kaseorg <anders@zulip.com>
Date: Wed, 30 Sep 2020 15:20:02 -0700
Subject: [PATCH] import_realm: Close the memcached connection before forking.

This prevents the memcached connection from being shared across
multiple processes, and hopefully addresses unexpected behavior from
cached functions like get_user_profile_by_id invoked inside the worker
processes.

Signed-off-by: Anders Kaseorg <anders@zulip.com>
---
 zerver/lib/import_realm.py | 2 ++
 zerver/lib/transfer.py     | 4 ++++
 2 files changed, 6 insertions(+)

diff --git a/zerver/lib/import_realm.py b/zerver/lib/import_realm.py
index c321069ad1..585148ca71 100644
--- a/zerver/lib/import_realm.py
+++ b/zerver/lib/import_realm.py
@@ -10,6 +10,7 @@ import boto3
 import orjson
 from bs4 import BeautifulSoup
 from django.conf import settings
+from django.core.cache import cache
 from django.db import connection
 from django.db.models import Max
 from django.utils.timezone import now as timezone_now
@@ -759,6 +760,7 @@ def import_uploads(realm: Realm, import_dir: Path, processes: int, processing_av
                 process_avatars(record)
         else:
             connection.close()
+            cache._cache.disconnect_all()
             with multiprocessing.Pool(processes) as p:
                 for out in p.imap_unordered(process_avatars, records):
                     pass
diff --git a/zerver/lib/transfer.py b/zerver/lib/transfer.py
index b9a9bcf71c..91c58a50d3 100644
--- a/zerver/lib/transfer.py
+++ b/zerver/lib/transfer.py
@@ -4,6 +4,7 @@ import os
 from mimetypes import guess_type
 
 from django.conf import settings
+from django.core.cache import cache
 from django.db import connection
 
 from zerver.lib.avatar_hash import user_avatar_path
@@ -35,6 +36,7 @@ def transfer_avatars_to_s3(processes: int) -> None:
             _transfer_avatar_to_s3(user)
     else:  # nocoverage
         connection.close()
+        cache._cache.disconnect_all()
         with multiprocessing.Pool(processes) as p:
             for out in p.imap_unordered(_transfer_avatar_to_s3, users):
                 pass
@@ -56,6 +58,7 @@ def transfer_message_files_to_s3(processes: int) -> None:
             _transfer_message_files_to_s3(attachment)
     else:  # nocoverage
         connection.close()
+        cache._cache.disconnect_all()
         with multiprocessing.Pool(processes) as p:
             for out in p.imap_unordered(_transfer_message_files_to_s3, attachments):
                 pass
@@ -82,6 +85,7 @@ def transfer_emoji_to_s3(processes: int) -> None:
             _transfer_emoji_to_s3(realm_emoji)
     else:  # nocoverage
         connection.close()
+        cache._cache.disconnect_all()
         with multiprocessing.Pool(processes) as p:
             for out in p.imap_unordered(_transfer_emoji_to_s3, realm_emojis):
                 pass