diff options
author | rowanbeentje <rowan@beent.je> | 2010-04-12 00:08:40 +0000 |
---|---|---|
committer | rowanbeentje <rowan@beent.je> | 2010-04-12 00:08:40 +0000 |
commit | ebfd8ca1dac81755451a22e364daa851992b386e (patch) | |
tree | ebc4ec3dbdbdb8e72f747d849a85a6a497bf6ea4 /Source/SPFileHandle.m | |
parent | cd9e9490ce6f073514cab18731cc0553f1c750d1 (diff) | |
download | sequelpro-ebfd8ca1dac81755451a22e364daa851992b386e.tar.gz sequelpro-ebfd8ca1dac81755451a22e364daa851992b386e.tar.bz2 sequelpro-ebfd8ca1dac81755451a22e364daa851992b386e.zip |
Add a new SPFileHandle class to support gzip compression and writing on a background thread, and integrate for SQL import:
- Implement streaming reading of gzip-compressed files for SQL import
- Support exporting SQL dumps into a gzip-compressed file
- SPFileHandle supports the most-used subset of NSFileHandle commands for easy integration
- Integrate zlib 1.2.4 for improved gzip streaming performance (and support for custom buffer sizes and file offset positions)
This implements Issue #571 .
Diffstat (limited to 'Source/SPFileHandle.m')
-rw-r--r-- | Source/SPFileHandle.m | 287 |
1 files changed, 287 insertions, 0 deletions
diff --git a/Source/SPFileHandle.m b/Source/SPFileHandle.m new file mode 100644 index 00000000..4b28d4d6 --- /dev/null +++ b/Source/SPFileHandle.m @@ -0,0 +1,287 @@ +// +// SPFileHandle.m +// sequel-pro +// +// Created by Rowan Beentje on 05/04/2010. +// Copyright 2010 Arboreal. All rights reserved. +// + +#import "SPFileHandle.h" +#import "zlib.1.2.4.h" + +// Define the size of the background read/write buffer. This affects speed and memory usage. +#define SPFH_BUFFER_SIZE 1048576 + +@interface SPFileHandle (PrivateAPI) +- (void) _writeBufferToData; +@end + + +@implementation SPFileHandle + +#pragma mark - +#pragma mark Setup and teardown + +/** + * Initialises and returns a SPFileHandle with a specified file (FILE or gzFile). + * "mode" indicates the file interaction mode - currently only read-only + * or write-only are supported. + * On reading, theFile should always be a gzFile; on writing, theFile is a FILE + * when compression is disabled, or a gzFile when enbled. + */ +- (id) initWithFile:(void *)theFile fromPath:(const char *)path mode:(int)mode +{ + if (self = [super init]) { + fileIsClosed = NO; + + wrappedFile = theFile; + wrappedFilePath = malloc(strlen(path) + 1); + strcpy(wrappedFilePath, path); + + // Check and set the mode + fileMode = mode; + if (fileMode != O_RDONLY && fileMode != O_WRONLY) { + [NSException raise:NSInvalidArgumentException format:@"SPFileHandle only supports read-only and write-only file modes"]; + } + + // Instantiate the buffer + pthread_mutex_init(&bufferLock, NULL); + buffer = [[NSMutableData alloc] init]; + bufferDataLength = 0; + bufferPosition = 0; + endOfFile = NO; + + // If in read mode, set up the buffer + if (fileMode == O_RDONLY) { + gzbuffer(wrappedFile, 131072); + useGzip = !gzdirect(wrappedFile); + processingThread = nil; + + // In write mode, set up a thread to handle writing in the background + } else if (fileMode == O_WRONLY) { + useGzip = NO; + processingThread = [[NSThread alloc] initWithTarget:self selector:@selector(_writeBufferToData) object:nil]; + [processingThread start]; + } + } + + return self; +} + +- (void) dealloc +{ + [self closeFile]; + if (processingThread) [processingThread release]; + free(wrappedFilePath); + [buffer release]; + pthread_mutex_destroy(&bufferLock); + [super dealloc]; +} + +#pragma mark - +#pragma mark Class methods + +/** + * Retrieve and return a SPFileHandle for reading a file at the supplied + * path. Returns nil if the file could not be found or opened. + */ ++ (id) fileHandleForReadingAtPath:(NSString *)path +{ + return [self fileHandleForPath:path mode:O_RDONLY]; +} + +/** + * Retrieve and return a SPFileHandle for writing a file at the supplied + * path. Returns nil if the file could not be found or opened. + */ ++ (id) fileHandleForWritingAtPath:(NSString *)path +{ + return [self fileHandleForPath:path mode:O_WRONLY]; +} + +/** + * Retrieve and return a SPFileHandle for a file at the specified path, + * using the supplied file status flag. Returns nil if the file could + * not be found or opened. + */ ++ (id) fileHandleForPath:(NSString *)path mode:(int)mode +{ + + // Retrieves the path in a filesystem-appropriate format and encoding + const char *pathRepresentation = [path fileSystemRepresentation]; + if (!pathRepresentation) return nil; + + // Open the file to get a file descriptor, returning on failure + FILE *theFile; + const char *theMode; + if (mode == O_WRONLY) { + theMode = "wb"; + theFile = fopen(pathRepresentation, theMode); + } else { + theMode = "rb"; + theFile = gzopen(pathRepresentation, theMode); + } + if (theFile == NULL) return nil; + + // Return an autoreleased file handle + return [[[self alloc] initWithFile:theFile fromPath:pathRepresentation mode:mode] autorelease]; +} + + +#pragma mark - +#pragma mark Data reading + +// Reads data up to a specified number of bytes from the file +- (NSMutableData *) readDataOfLength:(NSUInteger)length +{ + void *theData = malloc(length); + long theDataLength = gzread(wrappedFile, theData, length); + return [NSMutableData dataWithBytesNoCopy:theData length:theDataLength freeWhenDone:YES]; +} + +// Returns the data to the end of the file +- (NSMutableData *) readDataToEndOfFile +{ + return [self readDataOfLength:NSUIntegerMax]; +} + +// Returns the on-disk (raw) length of data read so far - can be used in progress bars +- (NSUInteger) realDataReadLength +{ + if (fileMode == O_WRONLY) return 0; + return gzoffset(wrappedFile); +} + +#pragma mark - +#pragma mark Data writing + +/** + * Set whether data should be written as gzipped data, defaulting + * to NO on a fresh object. If this is called after data has been + * written, an exception is thrown. + */ +- (void) setShouldWriteWithGzipCompression:(BOOL)shouldUseGzip +{ + if (shouldUseGzip == useGzip) return; + + if (dataWritten) [NSException raise:NSInternalInconsistencyException format:@"Cannot change compression settings when data has already been written"]; + + if (shouldUseGzip) { + fclose(wrappedFile); + wrappedFile = gzopen(wrappedFilePath, "wb"); + gzbuffer(wrappedFile, 131072); + } else { + gzclose(wrappedFile); + wrappedFile = fopen(wrappedFilePath, "wb"); + } + useGzip = shouldUseGzip; +} + + +// Write the provided data to the file +- (void) writeData:(NSData *)data +{ + + // Throw an exception if the file is closed + if (fileIsClosed) [NSException raise:NSInternalInconsistencyException format:@"Cannot write to a file handle after it has been closed"]; + + // Add the data to the buffer + pthread_mutex_lock(&bufferLock); + [buffer appendData:data]; + bufferDataLength += [data length]; + + // If the buffer is large, wait for some to be written out + while (bufferDataLength > SPFH_BUFFER_SIZE) { + pthread_mutex_unlock(&bufferLock); + usleep(100); + pthread_mutex_lock(&bufferLock); + } + pthread_mutex_unlock(&bufferLock); +} + +// Ensures any buffers are written to disk +- (void) synchronizeFile +{ + pthread_mutex_lock(&bufferLock); + while (bufferDataLength) { + pthread_mutex_unlock(&bufferLock); + usleep(100); + pthread_mutex_lock(&bufferLock); + } + pthread_mutex_unlock(&bufferLock); +} + +// Prevent further access to the file +- (void)closeFile +{ + if (!fileIsClosed) { + [self synchronizeFile]; + if (useGzip || fileMode == O_RDONLY) { + gzclose(wrappedFile); + } else { + fclose(wrappedFile); + } + if (processingThread) { + if ([processingThread isExecuting]) { + [processingThread cancel]; + while ([processingThread isExecuting]) usleep(100); + } + } + fileIsClosed = YES; + } +} + + +#pragma mark - +#pragma mark File information + +/** + * Returns whether gzip compression is enabled on the file. + */ +- (BOOL) isCompressed +{ + return useGzip; +} + +@end + +@implementation SPFileHandle (PrivateAPI) + +/** + * A method to be called on a background thread, allowing write data to build + * up in a buffer and write to disk in chunks as the buffer fills. This allows + * background compression of the data when using Gzip compression. + */ +- (void) _writeBufferToData +{ + NSAutoreleasePool *writePool = [[NSAutoreleasePool alloc] init]; + + // Process the buffer in a loop into the file, until cancelled + while (![processingThread isCancelled]) { + + // Check whether any data in the buffer needs to be written out - using thread locks for safety + pthread_mutex_lock(&bufferLock); + if (!bufferDataLength) { + pthread_mutex_unlock(&bufferLock); + usleep(1000); + continue; + } + + // Write out the data + long bufferLengthWrittenOut; + if (useGzip) { + bufferLengthWrittenOut = gzwrite(wrappedFile, [buffer bytes], bufferDataLength); + } else { + bufferLengthWrittenOut = fwrite([buffer bytes], 1, bufferDataLength, wrappedFile); + } + + // Update the buffer + CFDataDeleteBytes((CFMutableDataRef)buffer, CFRangeMake(0, bufferLengthWrittenOut)); + bufferDataLength = 0; + pthread_mutex_unlock(&bufferLock); + } + + [writePool drain]; +} + +@end |