opt(api): optimize update contention on the providers table (#24520)

The provider update handler was missing proper transaction handling,
causing quota deductions to be lost. This fix uses session.begin()
context manager for automatic commit/rollback, ensuring provider
quota updates are properly persisted.

Fixes #24356

api/core/tools/tool_manager.py

@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Literal, Optional, Union, cast
 
 import sqlalchemy as sa
 from pydantic import TypeAdapter
+from sqlalchemy.orm import Session
 from yarl import URL
 
 import contexts
@@ -617,8 +618,9 @@ class ToolManager:
                 WHERE tenant_id = :tenant_id
                 ORDER BY tenant_id, provider, is_default DESC, created_at DESC
                 """
-        ids = [row.id for row in db.session.execute(sa.text(sql), {"tenant_id": tenant_id}).all()]
-        return db.session.query(BuiltinToolProvider).where(BuiltinToolProvider.id.in_(ids)).all()
+        with Session(db.engine).no_autoflush as session:
+            ids = [row.id for row in session.execute(sa.text(sql), {"tenant_id": tenant_id}).all()]
+            return session.query(BuiltinToolProvider).where(BuiltinToolProvider.id.in_(ids)).all()
 
     @classmethod
     def list_providers_from_api(

api/events/event_handlers/update_provider_when_message_created.py

@@ -85,6 +85,7 @@ def handle(sender: Message, **kwargs):
         values=_ProviderUpdateValues(last_used=current_time),
         description="basic_last_used_update",
     )
+    logging.info("provider used, tenant_id=%s, provider_name=%s", tenant_id, provider_name)
     updates_to_perform.append(basic_update)
 
     # 2. Check if we need to deduct quota (system provider only)
@@ -186,9 +187,11 @@ def _execute_provider_updates(updates_to_perform: list[_ProviderUpdateOperation]
     if not updates_to_perform:
         return
 
+    updates_to_perform = sorted(updates_to_perform, key=lambda i: (i.filters.tenant_id, i.filters.provider_name))
+
     # Use SQLAlchemy's context manager for transaction management
     # This automatically handles commit/rollback
-    with Session(db.engine) as session:
+    with Session(db.engine) as session, session.begin():
         # Use a single transaction for all updates
         for update_operation in updates_to_perform:
             filters = update_operation.filters
@@ -212,10 +215,13 @@ def _execute_provider_updates(updates_to_perform: list[_ProviderUpdateOperation]
 
             # Prepare values dict for SQLAlchemy update
             update_values = {}
-            if values.last_used is not None:
-                update_values["last_used"] = values.last_used
+            # updateing to `last_used` is removed due to performance reason.
+            # ref: https://github.com/langgenius/dify/issues/24526
             if values.quota_used is not None:
                 update_values["quota_used"] = values.quota_used
+            # Skip the current update operation if no updates are required.
+            if not update_values:
+                continue
 
             # Build and execute the update statement
             stmt = update(Provider).where(*where_conditions).values(**update_values)

api/services/tools/builtin_tools_manage_service.py

@@ -546,54 +546,53 @@ class BuiltinToolManageService:
         # get all builtin providers
         provider_controllers = ToolManager.list_builtin_providers(tenant_id)
 
-        with db.session.no_autoflush:
-            # get all user added providers
-            db_providers: list[BuiltinToolProvider] = ToolManager.list_default_builtin_providers(tenant_id)
+        # get all user added providers
+        db_providers: list[BuiltinToolProvider] = ToolManager.list_default_builtin_providers(tenant_id)
 
-            # rewrite db_providers
-            for db_provider in db_providers:
-                db_provider.provider = str(ToolProviderID(db_provider.provider))
+        # rewrite db_providers
+        for db_provider in db_providers:
+            db_provider.provider = str(ToolProviderID(db_provider.provider))
 
-            # find provider
-            def find_provider(provider):
-                return next(filter(lambda db_provider: db_provider.provider == provider, db_providers), None)
+        # find provider
+        def find_provider(provider):
+            return next(filter(lambda db_provider: db_provider.provider == provider, db_providers), None)
 
-            result: list[ToolProviderApiEntity] = []
+        result: list[ToolProviderApiEntity] = []
 
-            for provider_controller in provider_controllers:
-                try:
-                    # handle include, exclude
-                    if is_filtered(
-                        include_set=dify_config.POSITION_TOOL_INCLUDES_SET,  # type: ignore
-                        exclude_set=dify_config.POSITION_TOOL_EXCLUDES_SET,  # type: ignore
-                        data=provider_controller,
-                        name_func=lambda x: x.identity.name,
-                    ):
-                        continue
+        for provider_controller in provider_controllers:
+            try:
+                # handle include, exclude
+                if is_filtered(
+                    include_set=dify_config.POSITION_TOOL_INCLUDES_SET,  # type: ignore
+                    exclude_set=dify_config.POSITION_TOOL_EXCLUDES_SET,  # type: ignore
+                    data=provider_controller,
+                    name_func=lambda x: x.identity.name,
+                ):
+                    continue
 
-                    # convert provider controller to user provider
-                    user_builtin_provider = ToolTransformService.builtin_provider_to_user_provider(
-                        provider_controller=provider_controller,
-                        db_provider=find_provider(provider_controller.entity.identity.name),
-                        decrypt_credentials=True,
+                # convert provider controller to user provider
+                user_builtin_provider = ToolTransformService.builtin_provider_to_user_provider(
+                    provider_controller=provider_controller,
+                    db_provider=find_provider(provider_controller.entity.identity.name),
+                    decrypt_credentials=True,
+                )
+
+                # add icon
+                ToolTransformService.repack_provider(tenant_id=tenant_id, provider=user_builtin_provider)
+
+                tools = provider_controller.get_tools()
+                for tool in tools or []:
+                    user_builtin_provider.tools.append(
+                        ToolTransformService.convert_tool_entity_to_api_entity(
+                            tenant_id=tenant_id,
+                            tool=tool,
+                            labels=ToolLabelManager.get_tool_labels(provider_controller),
+                        )
                     )
 
-                    # add icon
-                    ToolTransformService.repack_provider(tenant_id=tenant_id, provider=user_builtin_provider)
-
-                    tools = provider_controller.get_tools()
-                    for tool in tools or []:
-                        user_builtin_provider.tools.append(
-                            ToolTransformService.convert_tool_entity_to_api_entity(
-                                tenant_id=tenant_id,
-                                tool=tool,
-                                labels=ToolLabelManager.get_tool_labels(provider_controller),
-                            )
-                        )
-
-                    result.append(user_builtin_provider)
-                except Exception as e:
-                    raise e
+                result.append(user_builtin_provider)
+            except Exception as e:
+                raise e
 
         return BuiltinToolProviderSort.sort(result)
 
@@ -604,7 +603,7 @@ class BuiltinToolManageService:
         1.if the default provider exists, return the default provider
         2.if the default provider does not exist, return the oldest provider
         """
-        with Session(db.engine) as session:
+        with Session(db.engine).no_autoflush as session:
             try:
                 full_provider_name = provider_name
                 provider_id_entity = ToolProviderID(provider_name)

web/app/components/datasets/documents/list.tsx

@@ -7,7 +7,6 @@ import { pick, uniq } from 'lodash-es'
 import {
   RiArchive2Line,
   RiDeleteBinLine,
-  RiDownloadLine,
   RiEditLine,
   RiEqualizer2Line,
   RiLoopLeftLine,
@@ -35,7 +34,6 @@ import type { ColorMap, IndicatorProps } from '@/app/components/header/indicator
 import Indicator from '@/app/components/header/indicator'
 import { asyncRunSafe } from '@/utils'
 import { formatNumber } from '@/utils/format'
-import { useDocumentDownload } from '@/service/knowledge/use-document'
 import NotionIcon from '@/app/components/base/notion-icon'
 import ProgressBar from '@/app/components/base/progress-bar'
 import { ChunkingMode, DataSourceType, DocumentActionType, type DocumentDisplayStatus, type SimpleDocumentDetail } from '@/models/datasets'
@@ -189,7 +187,6 @@ export const OperationAction: FC<{
   scene?: 'list' | 'detail'
   className?: string
 }> = ({ embeddingAvailable, datasetId, detail, onUpdate, scene = 'list', className = '' }) => {
-  const downloadDocument = useDocumentDownload()
   const { id, enabled = false, archived = false, data_source_type, display_status } = detail || {}
   const [showModal, setShowModal] = useState(false)
   const [deleting, setDeleting] = useState(false)
@@ -298,31 +295,6 @@ export const OperationAction: FC<{
     )}
     {embeddingAvailable && (
       <>
-        <Tooltip
-          popupContent={t('datasetDocuments.list.action.download')}
-          popupClassName='text-text-secondary system-xs-medium'
-        >
-          <button
-            className={cn('mr-2 cursor-pointer rounded-lg',
-              !isListScene
-                ? 'shadow-shadow-3 border-[0.5px] border-components-button-secondary-border bg-components-button-secondary-bg p-2 shadow-xs backdrop-blur-[5px] hover:border-components-button-secondary-border-hover hover:bg-components-button-secondary-bg-hover'
-                : 'p-0.5 hover:bg-state-base-hover')}
-            onClick={() => {
-              downloadDocument.mutateAsync({
-                datasetId,
-                documentId: detail.id,
-                  }).then((response) => {
-                    if (response.download_url)
-                      window.location.href = response.download_url
-              }).catch((error) => {
-                console.error(error)
-                notify({ type: 'error', message: t('common.actionMsg.downloadFailed') })
-              })
-            }}
-          >
-            <RiDownloadLine className='h-4 w-4 text-components-button-secondary-text' />
-          </button>
-        </Tooltip>
         <Tooltip
           popupContent={t('datasetDocuments.list.action.settings')}
           popupClassName='text-text-secondary system-xs-medium'
@@ -526,7 +498,7 @@ const DocumentList: FC<IDocumentListProps> = ({
         const result = aValue.localeCompare(bValue)
         return sortOrder === 'asc' ? result : -result
       }
- else {
+      else {
         const result = aValue - bValue
         return sortOrder === 'asc' ? result : -result
       }
@@ -539,7 +511,7 @@ const DocumentList: FC<IDocumentListProps> = ({
     if (sortField === field) {
       setSortOrder(sortOrder === 'asc' ? 'desc' : 'asc')
     }
- else {
+    else {
       setSortField(field)
       setSortOrder('desc')
     }

web/i18n/en-US/dataset-documents.ts

@@ -32,7 +32,6 @@ const translation = {
       sync: 'Sync',
       pause: 'Pause',
       resume: 'Resume',
-      download: 'Download File',
     },
     index: {
       enable: 'Enable',

web/service/knowledge/use-document.ts

@@ -8,9 +8,7 @@ import type { MetadataType, SortType } from '../datasets'
 import { pauseDocIndexing, resumeDocIndexing } from '../datasets'
 import type { DocumentDetailResponse, DocumentListResponse, UpdateDocumentBatchParams } from '@/models/datasets'
 import { DocumentActionType } from '@/models/datasets'
-import type { CommonResponse, FileDownloadResponse } from '@/models/common'
-// Download document with authentication (sends Authorization header)
-import Toast from '@/app/components/base/toast'
+import type { CommonResponse } from '@/models/common'
 
 const NAME_SPACE = 'knowledge/document'
 
@@ -97,21 +95,6 @@ export const useSyncDocument = () => {
   })
 }
 
-// Download document with authentication (sends Authorization header)
-export const useDocumentDownload = () => {
-  return useMutation({
-    mutationFn: async ({ datasetId, documentId }: { datasetId: string; documentId: string }) => {
-      // The get helper automatically adds the Authorization header from localStorage
-      return get<FileDownloadResponse>(`/datasets/${datasetId}/documents/${documentId}/upload-file`)
-    },
-    onError: (error: any) => {
-      // Show a toast notification if download fails
-      const message = error?.message || 'Download failed.'
-      Toast.notify({ type: 'error', message })
-    },
-  })
-}
-
 export const useSyncWebsite = () => {
   return useMutation({
     mutationFn: ({ datasetId, documentId }: UpdateDocumentBatchParams) => {