feat: Calculate and display total problem file storage by summing all attached files instead of just test data.

Author: whlongg

judge/management/commands/backfill_problem_data_size.py

@@ -19,33 +19,51 @@ def handle(self, *args, **options):
         if dry_run:
             self.stdout.write(self.style.WARNING('DRY RUN MODE - No changes will be saved'))
 
-        problem_data_list = ProblemData.objects.select_related('problem').all()
-        total_count = problem_data_list.count()
+        from django.conf import settings
+        import os
+
+        problem_data_list = ProblemData.objects.all().iterator(chunk_size=5000)
+        total_count = ProblemData.objects.count()
         updated_count = 0
+        batch_size = 1000
+        batch = []
+        root = settings.DMOJ_PROBLEM_DATA_ROOT
 
         self.stdout.write(f'Processing {total_count} problems...\n')
 
         for problem_data in problem_data_list:
-            problem_code = problem_data.problem.code
             old_zipfile_size = problem_data.zipfile_size
+            total_size = 0
 
-            # Calculate new sizes
-            problem_data.update_zipfile_size()
+            # Calculate new sizes directly via OS to bypass Storage abstractions overhead
+            for field in ['zipfile', 'generator', 'custom_checker', 'custom_grader', 'custom_header']:
+                val = getattr(problem_data, field)
+                if val and val.name:
+                    path = os.path.join(root, val.name)
+                    try:
+                        total_size += os.path.getsize(path)
+                    except (OSError, FileNotFoundError):
+                        pass
 
-            new_zipfile_size = problem_data.zipfile_size
+            new_zipfile_size = total_size
 
             # Check if anything changed
             if old_zipfile_size != new_zipfile_size:
+                problem_data.zipfile_size = new_zipfile_size
+                batch.append(problem_data)
                 updated_count += 1
 
                 self.stdout.write(
-                    f'Problem: {problem_code}\n'
-                    f'  Test data: {self._format_size(old_zipfile_size)} -> {self._format_size(new_zipfile_size)}\n',
+                    f'Problem: {problem_data.problem_id}\n'
+                    f'  Total Storage: {self._format_size(old_zipfile_size)} -> {self._format_size(new_zipfile_size)}\n',
                 )
 
-                if not dry_run:
-                    # Use update_fields to avoid triggering save hooks again
-                    problem_data.save(update_fields=['zipfile_size'])
+                if len(batch) >= batch_size and not dry_run:
+                    ProblemData.objects.bulk_update(batch, ['zipfile_size'])
+                    batch = []
+
+        if batch and not dry_run:
+            ProblemData.objects.bulk_update(batch, ['zipfile_size'])
 
         if dry_run:
             self.stdout.write(self.style.WARNING(f'\nDRY RUN: Would update {updated_count}/{total_count} problems'))

judge/models/problem_data.py

@@ -99,15 +99,15 @@ def has_yml(self):
         return problem_data_storage.exists('%s/init.yml' % self.problem.code)
 
     def update_zipfile_size(self):
-        """Update the zipfile_size field based on the actual file size."""
-        if self.zipfile:
-            try:
-                self.zipfile_size = self.zipfile.size
-            except (OSError, IOError):
-                # If file doesn't exist or can't be accessed, set size to 0
-                self.zipfile_size = 0
-        else:
-            self.zipfile_size = 0
+        """Update the zipfile_size field based on the actual size of all attached files."""
+        total_size = 0
+        for field in [self.zipfile, self.generator, self.custom_checker, self.custom_grader, self.custom_header]:
+            if field:
+                try:
+                    total_size += field.size
+                except (OSError, IOError, ValueError):
+                    pass
+        self.zipfile_size = total_size
 
     def save(self, *args, **kwargs):
         # Update zipfile size before saving

judge/views/organization.py

@@ -789,25 +789,25 @@ def get_context_data(self, **kwargs):
             storage_totals = ProblemData.objects.filter(
                 problem__organization=self.organization,
             ).aggregate(
-                test_data=Sum('zipfile_size'),
+                total_storage=Sum('zipfile_size'),
             )
 
-            test_data_total = storage_totals['test_data'] or 0
+            total_storage_used = storage_totals['total_storage'] or 0
 
             cached_data = {
-                'test_data': test_data_total,
+                'total_storage': total_storage_used,
             }
 
             cache_factory.set_cache(cached_data)
 
-        context['test_data_storage'] = cached_data['test_data']
+        context['total_storage'] = cached_data.get('total_storage', cached_data.get('test_data', 0))
 
         # Quota information
         org = self.organization
         context['max_storage'] = org.get_max_storage()
         context['max_problems'] = org.get_max_problems()
         context['problem_count'] = org.get_current_problem_count()
-        context['storage_exceeded'] = context['test_data_storage'] > org.get_max_storage()
+        context['storage_exceeded'] = context['total_storage'] > org.get_max_storage()
         context['problem_limit_reached'] = context['problem_count'] >= org.get_max_problems()
 
         context.update(paginate_query_context(self.request))

templates/organization/storage.html

@@ -37,7 +37,7 @@
         <h3>{{ _('Storage Summary') }}</h3>
         <div style="margin-left: 2em; color: #666;">
             <p style="margin: 0.25em 0;">
-                {{ _('Test data:') }} {{ test_data_storage|filesizeformat }}
+                {{ _('Total Storage:') }} {{ total_storage|filesizeformat }}
                 / {{ max_storage|filesizeformat }}
             </p>
             <p style="margin: 0.25em 0;">
@@ -47,7 +47,7 @@ <h3>{{ _('Storage Summary') }}</h3>
 
         {% if storage_exceeded %}
         <div class="quota-warning">
-            {{ _('Storage limit reached! Please delete some test data to free up space.') }}
+            {{ _('Storage limit reached! Please delete some files or test data to free up space.') }}
         </div>
         {% endif %}
 
@@ -63,7 +63,7 @@ <h3>{{ _('Storage Summary') }}</h3>
             <tr>
                 <th>{{ _('Problem Code') }}</th>
                 <th>{{ _('Problem Name') }}</th>
-                <th>{{ _('Test Data') }}</th>
+                <th>{{ _('Total Size') }}</th>
             </tr>
         </thead>
         <tbody>