Skip to content

Commit

Permalink
STYLE: Replace thread_local variables with m_TensorStoreData data member
Browse files Browse the repository at this point in the history 
It appears preferable that multiple OMEZarrNGFFImageIO instances do not share
their TensorStore context and store. With this commit, each OMEZarrNGFFImageIO
instance will have its own context and store, stored in its `m_TensorStoreData`
data member.

The `m_TensorStoreData` data member hides the TensorStore data from the
interface, as it is a pointer to a forward-declared struct, similar to the
"Pimpl" idiom.
  • Loading branch information
@N-Dekker
N-Dekker committed Aug 27, 2024
1 parent 1d0f131 commit 3dd47f5
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 25 deletions.
4 changes: 4 additions & 0 deletions include/itkOMEZarrNGFFImageIO.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@


#include <fstream>
#include <memory> // For unique_ptr.
#include <string>
#include <vector>
#include "itkImageIOBase.h"
Expand Down Expand Up @@ -246,6 +247,9 @@ class IOOMEZarrNGFF_EXPORT OMEZarrNGFFImageIO : public ImageIOBase
char m_EmptyZip[m_EmptyZipSize] = "PK\x05\x06"; // the rest is filled with zeroes
const BufferInfo m_EmptyZipBufferInfo{ m_EmptyZip, m_EmptyZipSize };
const std::string m_EmptyZipFileName = MakeMemoryFileName(m_EmptyZipBufferInfo);

struct TensorStoreData;
const std::unique_ptr<TensorStoreData> m_TensorStoreData{};
};
} // end namespace itk

Expand Down
55 changes: 30 additions & 25 deletions src/itkOMEZarrNGFFImageIO.cxx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -374,9 +374,14 @@ MakeKVStoreHTTPDriverSpec(nlohmann::json & spec, const std::string & fullPath)

} // namespace

thread_local tensorstore::Context tsContext = tensorstore::Context::Default();
struct OMEZarrNGFFImageIO::TensorStoreData
{
tensorstore::Context tsContext{ tensorstore::Context::Default() };
tensorstore::TensorStore<> store{};
};

OMEZarrNGFFImageIO::OMEZarrNGFFImageIO()
: m_TensorStoreData(std::make_unique<TensorStoreData>())
{
this->AddSupportedWriteExtension(".zarr");
this->AddSupportedWriteExtension(".zr2");
Expand Down Expand Up @@ -407,7 +412,7 @@ OMEZarrNGFFImageIO::PrintSelf(std::ostream & os, Indent indent) const

// JSON file path, e.g. "C:/Dev/ITKIOOMEZarrNGFF/v0.4/cyx.ome.zarr/.zgroup"
void
writeJson(nlohmann::json json, std::string path, std::string driver)
writeJson(nlohmann::json json, std::string path, std::string driver, tensorstore::Context& tsContext)
{
auto attrs_store = tensorstore::Open<nlohmann::json, 0>(
{ { "driver", "json" }, { "kvstore", { { "driver", driver }, { "path", path } } } },
Expand All @@ -428,7 +433,7 @@ writeJson(nlohmann::json json, std::string path, std::string driver)

// JSON file path, e.g. "C:/Dev/ITKIOOMEZarrNGFF/v0.4/cyx.ome.zarr/.zattrs"
bool
jsonRead(const std::string path, nlohmann::json & result, std::string driver)
jsonRead(const std::string path, nlohmann::json & result, std::string driver, tensorstore::Context& tsContext)
{
// Reading JSON via TensorStore allows it to be in the cloud
nlohmann::json readSpec = { { "driver", "json" }, { "kvstore", { { "driver", driver }, { "path", path } } } };
Expand Down Expand Up @@ -466,15 +471,15 @@ OMEZarrNGFFImageIO::CanReadFile(const char * filename)
{
std::string driver = getKVstoreDriver(filename);
nlohmann::json json;
if (!jsonRead(std::string(filename) + "/.zgroup", json, driver))
if (!jsonRead(std::string(filename) + "/.zgroup", json, driver, m_TensorStoreData->tsContext))
{
return false;
}
if (json.at("zarr_format").get<int>() != 2)
{
return false; // unsupported zarr format
}
if (!jsonRead(std::string(filename) + "/.zattrs", json, driver))
if (!jsonRead(std::string(filename) + "/.zattrs", json, driver, m_TensorStoreData->tsContext))
{
return false;
}
Expand All @@ -491,8 +496,6 @@ OMEZarrNGFFImageIO::CanReadFile(const char * filename)
// return this->HasSupportedWriteExtension(filename, true);
}

thread_local tensorstore::TensorStore<> store; // initialized by ReadImageInformation/ReadArrayMetadata

void
OMEZarrNGFFImageIO::ReadArrayMetadata(std::string path, std::string driver)
{
Expand All @@ -503,15 +506,15 @@ OMEZarrNGFFImageIO::ReadArrayMetadata(std::string path, std::string driver)
}

auto openFuture = tensorstore::Open(readSpec,
tsContext,
m_TensorStoreData->tsContext,
tensorstore::OpenMode::open,
tensorstore::RecheckCached{ false },
tensorstore::ReadWriteMode::read);
TS_EVAL_CHECK(openFuture);
store = openFuture.value();
auto shape_span = store.domain().shape();
m_TensorStoreData->store = openFuture.value();
auto shape_span = m_TensorStoreData->store.domain().shape();

tensorstore::DataType dtype = store.dtype();
tensorstore::DataType dtype = m_TensorStoreData->store.dtype();
this->SetComponentType(tensorstoreToITKComponentType(dtype));

std::vector<int64_t> dims(shape_span.rbegin(), shape_span.rend()); // convert KJI into IJK
Expand All @@ -534,13 +537,15 @@ OMEZarrNGFFImageIO::ReadArrayMetadata(std::string path, std::string driver)
ImageIORegion
OMEZarrNGFFImageIO::ConfigureTensorstoreIORegion(const ImageIORegion & ioRegion) const
{
const auto storeRank = m_TensorStoreData->store.rank();

// Set up IO region to match known store dimensions
itkAssertOrThrowMacro(m_StoreAxes.size() == store.rank(), "Detected mismatch in axis count and store rank");
ImageIORegion storeRegion(store.rank());
itkAssertOrThrowMacro(storeRegion.GetImageDimension(), store.rank());
itkAssertOrThrowMacro(m_StoreAxes.size() == storeRank, "Detected mismatch in axis count and store rank");
ImageIORegion storeRegion(storeRank);
itkAssertOrThrowMacro(storeRegion.GetImageDimension(), storeRank);
auto storeAxes = this->GetAxesInStoreOrder();

for (size_t storeIndex = 0; storeIndex < store.rank(); ++storeIndex)
for (size_t storeIndex = 0; storeIndex < storeRank; ++storeIndex)
{
auto axisName = storeAxes.at(storeIndex).name;

Expand Down Expand Up @@ -651,12 +656,12 @@ OMEZarrNGFFImageIO::ReadImageInformation()
std::string driver = getKVstoreDriver(this->GetFileName());

const std::string zgroupFilePath(std::string(this->GetFileName()) + "/.zgroup");
bool status = jsonRead(zgroupFilePath, json, driver);
bool status = jsonRead(zgroupFilePath, json, driver, m_TensorStoreData->tsContext);
itkAssertOrThrowMacro(status, ("Failed to read from " + zgroupFilePath));
itkAssertOrThrowMacro(json.at("zarr_format").get<int>() == 2, "Only v2 zarr format is supported"); // only v2 for now

const std::string zattrsFilePath(std::string(this->GetFileName()) + "/.zattrs");
status = jsonRead(zattrsFilePath, json, driver);
status = jsonRead(zattrsFilePath, json, driver, m_TensorStoreData->tsContext);
itkAssertOrThrowMacro(status, ("Failed to read from " + zattrsFilePath));
json = json.at("multiscales")[0]; // multiscales must be present in OME-NGFF
auto version = json.at("version").get<std::string>();
Expand Down Expand Up @@ -732,7 +737,7 @@ OMEZarrNGFFImageIO::Read(void * buffer)
// This comparison needs to be done carefully, we can compare 3D and 6D regions
if (this->GetLargestRegion().GetNumberOfPixels() == m_IORegion.GetNumberOfPixels())
{
itkAssertOrThrowMacro(store.domain().num_elements() == m_IORegion.GetNumberOfPixels(),
itkAssertOrThrowMacro(m_TensorStoreData->store.domain().num_elements() == m_IORegion.GetNumberOfPixels(),
"Detected mismatch between store size and size of largest possible region");
}
else
Expand All @@ -750,7 +755,7 @@ OMEZarrNGFFImageIO::Read(void * buffer)
}

if (const IOComponentEnum componentType{ this->GetComponentType() };
!TryToReadFromStore(supportedPixelTypes, componentType, store, storeIORegion, buffer))
!TryToReadFromStore(supportedPixelTypes, componentType, m_TensorStoreData->store, storeIORegion, buffer))
{
itkExceptionMacro("Unsupported component type: " << GetComponentTypeAsString(componentType));
}
Expand All @@ -771,7 +776,7 @@ OMEZarrNGFFImageIO::WriteImageInformation()

nlohmann::json group;
group["zarr_format"] = 2;
writeJson(group, std::string(this->GetFileName()) + "/.zgroup", driver);
writeJson(group, std::string(this->GetFileName()) + "/.zgroup", driver, m_TensorStoreData->tsContext);

unsigned dim = this->GetNumberOfDimensions();

Expand Down Expand Up @@ -803,7 +808,7 @@ OMEZarrNGFFImageIO::WriteImageInformation()

nlohmann::json zattrs;
zattrs["multiscales"] = multiscales;
writeJson(zattrs, std::string(this->GetFileName()) + "/.zattrs", driver);
writeJson(zattrs, std::string(this->GetFileName()) + "/.zattrs", driver, m_TensorStoreData->tsContext);
}


Expand All @@ -812,7 +817,7 @@ OMEZarrNGFFImageIO::Write(const void * buffer)
{
if (m_FileName.substr(m_FileName.size() - 4) == ".zip" || m_FileName.substr(m_FileName.size() - 7) == ".memory")
{
tsContext = tensorstore::Context::Default(); // start with clean zip handles
m_TensorStoreData->tsContext = tensorstore::Context::Default(); // start with clean zip handles
}
this->WriteImageInformation();

Expand All @@ -837,8 +842,8 @@ OMEZarrNGFFImageIO::Write(const void * buffer)

if (!TryToWriteToStore(supportedPixelTypes,
componentType,
store,
tsContext,
m_TensorStoreData->store,
m_TensorStoreData->tsContext,
m_FileName,
MakePath(this->GetDatasetIndex()),
shape,
Expand All @@ -851,7 +856,7 @@ OMEZarrNGFFImageIO::Write(const void * buffer)
{
// Attempt to read a non-existent file from the in-memory zip to close the current one
nlohmann::json temp;
bool wasRead = jsonRead(m_EmptyZipFileName + "/non-existent.json", temp, "zip_memory");
bool wasRead = jsonRead(m_EmptyZipFileName + "/non-existent.json", temp, "zip_memory", m_TensorStoreData->tsContext);
assert(wasRead == false);
}
}
Expand Down

0 comments on commit 3dd47f5

Please sign in to comment.