persistent DuckDB connection to avoid Windows heap corruption

csaybar · csaybar · commit 6ff465c3e1f3 · 2026-05-21T16:22:51.000-05:00
diff --git a/.github/workflows/julia_fix.yml b/.github/workflows/julia_fix.yml
@@ -0,0 +1,142 @@
+name: Julia Fix CI
+
+# Scoped to the julia_fix branch so it doesn't fire on every push.
+# Workflow_dispatch lets you replay it manually after a force-push.
+on:
+  push:
+    branches: [julia_fix]
+  pull_request:
+    branches: [julia_fix]
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+
+jobs:
+  build-lib:
+    name: libcozip (${{ matrix.name }})
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - name: linux-x86_64
+            runner: ubuntu-24.04
+            container: quay.io/pypa/manylinux_2_28_x86_64
+            lib_ext: so
+            shell: bash
+          - name: macos-universal
+            runner: macos-14
+            container: ''
+            lib_ext: dylib
+            shell: bash
+            cmake_extra: >-
+              "-DCMAKE_OSX_ARCHITECTURES=arm64;x86_64"
+              -DCMAKE_OSX_DEPLOYMENT_TARGET=11.0
+          - name: windows-x86_64
+            runner: windows-2022
+            container: ''
+            lib_ext: dll
+            shell: pwsh
+    runs-on: ${{ matrix.runner }}
+    container: ${{ matrix.container }}
+    defaults:
+      run:
+        shell: ${{ matrix.shell }}
+    steps:
+      - uses: actions/checkout@v4
+
+      # manylinux ships gcc/cmake but not ninja.
+      - if: startsWith(matrix.name, 'linux')
+        run: yum install -y ninja-build
+      - if: runner.os == 'macOS'
+        run: brew install ninja
+      - if: runner.os == 'Windows'
+        run: choco install ninja -y
+
+      # cl.exe + nmake on PATH for cmake.
+      - if: runner.os == 'Windows'
+        uses: ilammy/msvc-dev-cmd@v1
+        with:
+          arch: x64
+
+      - name: Configure
+        run: >
+          cmake -B build -S core -G Ninja
+          -DCMAKE_BUILD_TYPE=Release
+          -DCMAKE_INSTALL_PREFIX=install
+          -DCMAKE_INSTALL_LIBDIR=lib
+          -DBUILD_SHARED_LIBS=ON
+          ${{ matrix.cmake_extra }}
+
+      - run: cmake --build build --config Release --parallel
+      - run: cmake --install build --config Release
+
+      # Julia's _resolve_lib_path expects the shared lib at the root of
+      # the artifact (COZIP_LIB_PATH points straight at it). No header
+      # or license needed for the test step.
+      - name: Stage lib (Unix)
+        if: runner.os != 'Windows'
+        run: |
+          mkdir -p staged
+          cp install/lib/cozip.${{ matrix.lib_ext }} staged/
+
+      - name: Stage lib (Windows)
+        if: runner.os == 'Windows'
+        run: |
+          New-Item -ItemType Directory -Force -Path staged | Out-Null
+          Copy-Item install\bin\cozip.dll staged\
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: libcozip-${{ matrix.name }}
+          path: staged/
+          if-no-files-found: error
+
+
+  build-julia-package:
+    name: Julia package (${{ matrix.name }})
+    needs: build-lib
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - name: linux-x86_64
+            runner: ubuntu-24.04
+            lib_ext: so
+          - name: macos-universal
+            runner: macos-14
+            lib_ext: dylib
+          - name: windows-x86_64
+            runner: windows-2022
+            lib_ext: dll
+    runs-on: ${{ matrix.runner }}
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: julia-actions/setup-julia@v2
+        with:
+          version: '1'
+
+      - uses: julia-actions/cache@v2
+
+      - uses: actions/download-artifact@v4
+        with:
+          name: libcozip-${{ matrix.name }}
+          path: libcozip-bin/
+
+      # Override Artifacts.toml so the test uses the just-built lib,
+      # not whatever the last published release points at.
+      - name: Set COZIP_LIB_PATH (Unix)
+        if: runner.os != 'Windows'
+        run: echo "COZIP_LIB_PATH=$(pwd)/libcozip-bin/cozip.${{ matrix.lib_ext }}" >> "$GITHUB_ENV"
+
+      - name: Set COZIP_LIB_PATH (Windows)
+        if: runner.os == 'Windows'
+        shell: pwsh
+        run: |
+          $p = Join-Path (Get-Location) "libcozip-bin\cozip.dll"
+          Add-Content -Path $env:GITHUB_ENV -Value "COZIP_LIB_PATH=$p"
+
+      - name: Test
+        run: julia --project=julia -e 'using Pkg; Pkg.instantiate(); Pkg.test()'
diff --git a/julia/src/Cozip.jl b/julia/src/Cozip.jl
@@ -6,6 +6,13 @@ include("Reader.jl")
 
 const write = create
 
+# Open the shared DuckDB connection once at module load and register
+# an atexit hook to close it before Julia unloads DLLs.
+function __init__()
+    _init_duckdb!()
+    atexit(_close_duckdb!)
+end
+
 export stage_metadata, stage_create, create
 
 end # module
diff --git a/julia/src/Reader.jl b/julia/src/Reader.jl
@@ -6,6 +6,26 @@ using DuckDB
 function read end
 
 
+# Cozip and httpfs are loaded lazily on the first Reader.read call,
+# not in Cozip.__init__, so `using Cozip` stays fully offline. The
+# writer path doesn't need either extension and shouldn't pay for
+# the community-registry roundtrip on cold cache.
+const _DUCKDB_EXTS_LOADED = Ref(false)
+
+
+# Always called while holding _DUCKDB_LOCK (from Writer.jl), so the
+# check-then-set is race-free without extra locking.
+function _ensure_duckdb_extensions!(con)
+    _DUCKDB_EXTS_LOADED[] && return nothing
+    DBInterface.execute(con, "INSTALL httpfs")
+    DBInterface.execute(con, "LOAD httpfs")
+    DBInterface.execute(con, "INSTALL cozip FROM community")
+    DBInterface.execute(con, "LOAD cozip")
+    _DUCKDB_EXTS_LOADED[] = true
+    return nothing
+end
+
+
 """
     read(source; columns=nothing, gdal_vsi=true) -> DataFrame
 
@@ -33,22 +53,16 @@ function read(
         throw(ArgumentError("cozip: `columns` entries must be non-empty"))
     end
 
-    db  = DuckDB.DB()
-    con = DBInterface.connect(db)
-    try
-        DBInterface.execute(con, "INSTALL httpfs")
-        DBInterface.execute(con, "LOAD httpfs")
-        DBInterface.execute(con, "INSTALL cozip FROM community")
-        DBInterface.execute(con, "LOAD cozip")
+    con = _duckdb_con()
+
+    return lock(_DUCKDB_LOCK) do
+        _ensure_duckdb_extensions!(con)
 
         sql = string(
             "SELECT ", _build_select(columns, gdal_vsi),
             " FROM read_cozip(?, gdal_vsi := ", gdal_vsi ? "true" : "false", ")",
         )
-        return DataFrame(DBInterface.execute(con, sql, [src]))
-    finally
-        DBInterface.close!(con)
-        DBInterface.close!(db)
+        DataFrame(DBInterface.execute(con, sql, [src]))
     end
 end
 
diff --git a/julia/src/Writer.jl b/julia/src/Writer.jl
@@ -18,6 +18,49 @@ const _RESERVED_INPUT_COLUMNS    = Set(["offset", "size"])
 const _REQUIRED_METADATA_COLUMNS = Set(["name", "offset", "size"])
 
 
+# Shared DuckDB connection used by every parquet helper here and by
+# Reader.read.
+
+const _DUCKDB_DB   = Ref{Union{Nothing,DuckDB.DB}}(nothing)
+const _DUCKDB_CON  = Ref{Union{Nothing,DuckDB.Connection}}(nothing)
+const _DUCKDB_LOCK = ReentrantLock()
+
+
+function _init_duckdb!()
+    _DUCKDB_CON[] === nothing || return nothing
+    _DUCKDB_DB[]  = DuckDB.DB()
+    _DUCKDB_CON[] = DBInterface.connect(_DUCKDB_DB[])
+    return nothing
+end
+
+
+function _close_duckdb!()
+    if _DUCKDB_CON[] !== nothing
+        try
+            DBInterface.close!(_DUCKDB_CON[])
+        catch
+        end
+        _DUCKDB_CON[] = nothing
+    end
+    if _DUCKDB_DB[] !== nothing
+        try
+            DBInterface.close!(_DUCKDB_DB[])
+        catch
+        end
+        _DUCKDB_DB[] = nothing
+    end
+    return nothing
+end
+
+
+function _duckdb_con()
+    con = _DUCKDB_CON[]
+    con === nothing &&
+        error("cozip: DuckDB connection not initialized; Cozip.__init__ must run first")
+    return con
+end
+
+
 """
     stage_metadata(table) -> (metadata, paths)
 
@@ -305,31 +348,25 @@ function _alloc_user_entries(
 end
 
 
-# Cached: short-lived connections trigger Windows heap corruption (duckdb#10787).
-const _DUCKDB_DB  = Ref{Any}(nothing)
-const _DUCKDB_CON = Ref{Any}(nothing)
-
-function _duckdb_con()
-    if _DUCKDB_CON[] === nothing
-        _DUCKDB_DB[]  = DuckDB.DB()
-        _DUCKDB_CON[] = DBInterface.connect(_DUCKDB_DB[])
+function _read_parquet_columns(path::String)::Vector{String}
+    con = _duckdb_con()
+    return lock(_DUCKDB_LOCK) do
+        result = DBInterface.execute(con,
+            "SELECT * FROM read_parquet('$(_sql_escape(path))') LIMIT 0")
+        names(DataFrame(result))
     end
-    return _DUCKDB_CON[]
 end
 
 
-function _read_parquet_columns(path::String)::Vector{String}
-    result = DBInterface.execute(_duckdb_con(),
-        "SELECT * FROM read_parquet('$(_sql_escape(path))') LIMIT 0")
-    return names(DataFrame(result))
-end
-
 function _read_parquet_subset(path::String, columns::Vector{String})::DataFrame
-    # `offset` and `size` are reserved in DuckDB SQL.
-    cols_sql = join(["\"$c\"" for c in columns], ", ")
-    result = DBInterface.execute(_duckdb_con(),
-        "SELECT $cols_sql FROM read_parquet('$(_sql_escape(path))')")
-    return DataFrame(result)
+    con = _duckdb_con()
+    return lock(_DUCKDB_LOCK) do
+        # `offset` and `size` are reserved in DuckDB SQL.
+        cols_sql = join(["\"$c\"" for c in columns], ", ")
+        result = DBInterface.execute(con,
+            "SELECT $cols_sql FROM read_parquet('$(_sql_escape(path))')")
+        DataFrame(result)
+    end
 end
 
 
@@ -390,22 +427,27 @@ function _write_metadata_parquet(df::DataFrame, out_path::AbstractString)
     cols      = names(df)
     col_types = [string("\"", c, "\" ", _julia_to_duckdb_type(eltype(df[!, c])))
                  for c in cols]
-    tbl = "tmp_cozip_$(getpid())_$(time_ns())"
+
     con = _duckdb_con()
 
-    try
-        DBInterface.execute(con,
-            "CREATE TABLE $tbl ($(join(col_types, ", ")))")
-        for i in 1:nrow(df)
-            vals = join((_sql_literal(df[i, c]) for c in cols), ", ")
-            DBInterface.execute(con, "INSERT INTO $tbl VALUES ($vals)")
+    lock(_DUCKDB_LOCK) do
+        try
+            DBInterface.execute(con,
+                "CREATE OR REPLACE TEMP TABLE tmp_cozip ($(join(col_types, ", ")))")
+
+            for i in 1:nrow(df)
+                vals = join((_sql_literal(df[i, c]) for c in cols), ", ")
+                DBInterface.execute(con, "INSERT INTO tmp_cozip VALUES ($vals)")
+            end
+
+            kv_sql = _kv_metadata_sql(_table_metadata(df))
+            DBInterface.execute(con,
+                "COPY tmp_cozip TO '$(_sql_escape(String(out_path)))' (FORMAT parquet$kv_sql)")
+        finally
+            DBInterface.execute(con, "DROP TABLE IF EXISTS tmp_cozip")
         end
-        kv_sql = _kv_metadata_sql(_table_metadata(df))
-        DBInterface.execute(con,
-            "COPY $tbl TO '$(_sql_escape(String(out_path)))' (FORMAT parquet$kv_sql)")
-    finally
-        DBInterface.execute(con, "DROP TABLE IF EXISTS $tbl")
     end
+
     return nothing
 end
 
diff --git a/julia/test/runtests.jl b/julia/test/runtests.jl
