Skip to content

fix: Address multiple memory leaks in CodeBlock component #4244

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jun 23, 2025
Merged

fix: Address multiple memory leaks in CodeBlock component #4244

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4313d4c
fix: address multiple memory leaks in CodeBlock component (CodeBlock_…
kiwina Jun 2, 2025
1e30122
fix: Address review feedback for CodeBlock memory leak fixes
daniel-lxs Jun 23, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 68 additions & 16 deletions webview-ui/src/components/common/CodeBlock.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -232,6 +232,10 @@ const CodeBlock = memo(
const copyButtonWrapperRef = useRef<HTMLDivElement>(null)
const { showCopyFeedback, copyWithFeedback } = useCopyToClipboard()
const { t } = useAppTranslation()
const isMountedRef = useRef(true)
const buttonPositionTimeoutRef = useRef<NodeJS.Timeout | null>(null)
const collapseTimeout1Ref = useRef<NodeJS.Timeout | null>(null)
const collapseTimeout2Ref = useRef<NodeJS.Timeout | null>(null)

// Update current language when prop changes, but only if user hasn't
// made a selection.
Expand All @@ -243,17 +247,40 @@ const CodeBlock = memo(
}
}, [language, currentLanguage])

// Manage mounted state
useEffect(() => {
isMountedRef.current = true
return () => {
isMountedRef.current = false
}
}, [])
Copy link
Member

@daniel-lxs daniel-lxs Jun 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a specific reason for having a dedicated useEffect just for managing isMountedRef? This could potentially be integrated into the existing syntax highlighting effect since they have the same lifecycle. Just curious about the design choice here!


// Cleanup for collapse/expand timeouts
useEffect(() => {
return () => {
if (collapseTimeout1Ref.current) {
clearTimeout(collapseTimeout1Ref.current)
}
if (collapseTimeout2Ref.current) {
clearTimeout(collapseTimeout2Ref.current)
}
}
}, [])
Copy link
Member

@daniel-lxs daniel-lxs Jun 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I noticed that the cleanup pattern for timeouts is a bit inconsistent. The buttonPositionTimeoutRef is cleaned up in the same effect that sets it (lines 521-526), while the collapse timeouts have a separate cleanup effect here.

Could we consider consolidating these patterns? Having a consistent approach would improve maintainability. What do you think about either:

  1. Moving all timeout cleanups to their respective effects, or
  2. Having a single cleanup effect for all timeouts?

Both approaches would work, just wondering if there's a preference!


// Syntax highlighting with cached Shiki instance.
useEffect(() => {
const fallback = `<pre style="padding: 0; margin: 0;"><code class="hljs language-${currentLanguage || "txt"}">${source || ""}</code></pre>`

const highlight = async () => {
// Show plain text if language needs to be loaded.
if (currentLanguage && !isLanguageLoaded(currentLanguage)) {
setHighlightedCode(fallback)
if (isMountedRef.current) {
setHighlightedCode(fallback)
}
}

const highlighter = await getHighlighter(currentLanguage)
if (!isMountedRef.current) return

const html = await highlighter.codeToHtml(source || "", {
lang: currentLanguage || "txt",
Expand All @@ -277,13 +304,18 @@ const CodeBlock = memo(
},
] as ShikiTransformer[],
})
if (!isMountedRef.current) return

setHighlightedCode(html)
if (isMountedRef.current) {
setHighlightedCode(html)
}
}

highlight().catch((e) => {
console.error("[CodeBlock] Syntax highlighting error:", e, "\nStack trace:", e.stack)
setHighlightedCode(fallback)
if (isMountedRef.current) {
setHighlightedCode(fallback)
}
})
}, [source, currentLanguage, collapsedHeight])

Expand Down Expand Up @@ -455,8 +487,15 @@ const CodeBlock = memo(
// Update button position and scroll when highlightedCode changes
useEffect(() => {
if (highlightedCode) {
// Clear any existing timeout before setting a new one
if (buttonPositionTimeoutRef.current) {
clearTimeout(buttonPositionTimeoutRef.current)
}
// Update button position
setTimeout(updateCodeBlockButtonPosition, 0)
buttonPositionTimeoutRef.current = setTimeout(() => {
updateCodeBlockButtonPosition()
buttonPositionTimeoutRef.current = null // Optional: Clear ref after execution
}, 0)

// Scroll to bottom if needed (immediately after Shiki updates)
if (shouldScrollAfterHighlightRef.current) {
Expand All @@ -479,6 +518,12 @@ const CodeBlock = memo(
shouldScrollAfterHighlightRef.current = false
}
}
// Cleanup function for this effect
return () => {
if (buttonPositionTimeoutRef.current) {
clearTimeout(buttonPositionTimeoutRef.current)
}
}
}, [highlightedCode, updateCodeBlockButtonPosition])

// Advanced inertial scroll chaining
Expand Down Expand Up @@ -682,23 +727,30 @@ const CodeBlock = memo(
{showCollapseButton && (
<CodeBlockButton
onClick={() => {
// Get the current code block element and scrollable container
const codeBlock = codeBlockRef.current
const scrollContainer = document.querySelector('[data-virtuoso-scroller="true"]')
if (!codeBlock || !scrollContainer) return

// Get the current code block element
const codeBlock = codeBlockRef.current // Capture ref early
// Toggle window shade state
setWindowShade(!windowShade)

// Clear any previous timeouts
if (collapseTimeout1Ref.current) clearTimeout(collapseTimeout1Ref.current)
if (collapseTimeout2Ref.current) clearTimeout(collapseTimeout2Ref.current)

// After UI updates, ensure code block is visible and update button position
setTimeout(
collapseTimeout1Ref.current = setTimeout(
() => {
codeBlock.scrollIntoView({ behavior: "smooth", block: "nearest" })

// Wait for scroll to complete before updating button position
setTimeout(() => {
updateCodeBlockButtonPosition()
}, 50)
if (codeBlock) {
// Check if codeBlock element still exists
codeBlock.scrollIntoView({ behavior: "smooth", block: "nearest" })

// Wait for scroll to complete before updating button position
collapseTimeout2Ref.current = setTimeout(() => {
// updateCodeBlockButtonPosition itself should also check for refs if needed
updateCodeBlockButtonPosition()
collapseTimeout2Ref.current = null
}, 50)
}
collapseTimeout1Ref.current = null
},
WINDOW_SHADE_SETTINGS.transitionDelayS * 1000 + 50,
)
Expand Down