Associate Handle Addition

include/sys/zia_cddl.h

@@ -32,7 +32,7 @@
 int
 zia_compress_impl(const dpusm_uf_t *dpusm, zia_props_t *props,
     enum zio_compress c, abd_t *src, size_t s_len,
-    void **cbuf_handle, uint64_t *c_len,
+    abd_t **dst, void **cbuf_handle, uint64_t *c_len,
     uint8_t level, boolean_t *local_offload);
 
 int

module/zfs/spa.c

@@ -901,7 +901,8 @@ spa_prop_validate(spa_t *spa, nvlist_t *props)
 
 			void *new_provider = zia_get_provider(strval);
 			if (new_provider == NULL) {
-				error = SET_ERROR(ZFS_ERR_ZIA_NONEXISTENT_PROVIDER);
+				error = SET_ERROR(
+				    ZFS_ERR_ZIA_NONEXISTENT_PROVIDER);
 				break;
 			}
 

module/zfs/zia.c

@@ -1001,12 +1001,14 @@ zia_compress(zia_props_t *props, enum zio_compress c,
 	}
 
 	void *cbuf_handle = NULL;
+	*dst = abd_alloc_sametype(src, s_len);
 	const int rc = zia_compress_impl(dpusm, props, c, src, s_len,
-	    &cbuf_handle, d_len, level, local_offload);
-	if (rc == ZIA_OK) {
-		*dst = abd_alloc_sametype(src, s_len);
-		ABD_HANDLE(*dst) = cbuf_handle;
+	    dst, &cbuf_handle, d_len, level, local_offload);
+	if (rc != ZIA_OK) {
+		abd_free(*dst);
+		return (rc);
 	}
+	ABD_HANDLE(*dst) = cbuf_handle;
 	return (rc);
 #else
 	(void) props; (void) c; (void) src; (void) s_len;
@@ -1048,11 +1050,13 @@ zia_decompress(zia_props_t *props, enum zio_compress c,
 	 *
 	 * a lot of these will fail because d_len tends to be small
 	 */
-	ABD_HANDLE(dst) = zia_alloc(props->provider, d_len);
+	void *handle = zia_alloc(props->provider, d_len);
+	ABD_HANDLE(dst) = handle;
 	if (!ABD_HANDLE(dst)) {
 		/* let abd_free clean up zio->io_abd */
 		return (ZIA_ERROR);
 	}
+	dpusm->associate_handle(handle, ABD_LINEAR_BUF(dst));
 
 	/*
 	 * d_len pulled from accelerator is not used, so
@@ -1328,6 +1332,8 @@ zia_raidz_alloc(zio_t *zio, raidz_row_t *rr, boolean_t rec,
 			goto error;
 		}
 
+		dpusm->associate_handle(handle, ABD_LINEAR_BUF(rc->rc_abd));
+
 		if (dpusm->raid.set_column(rr->rr_zia_handle,
 		    c, handle, rc->rc_size) != DPUSM_OK) {
 			goto error;
@@ -1511,6 +1517,9 @@ zia_raidz_new_parity(zio_t *zio, raidz_row_t *rr, uint64_t c)
 		return (ZIA_ERROR);
 	}
 
+	dpusm->associate_handle(new_parity_handle,
+	    ABD_LINEAR_BUF(rc->rc_abd));
+
 	const int ret = dpusm->raid.set_column(rr->rr_zia_handle,
 	    c, new_parity_handle, rc->rc_size);
 	if (ret == DPUSM_OK) {

module/zfs/zia_cddl.c

@@ -34,7 +34,7 @@
 int
 zia_compress_impl(const dpusm_uf_t *dpusm, zia_props_t *props,
     enum zio_compress c, abd_t *src, size_t s_len,
-    void **cbuf_handle, uint64_t *c_len,
+    abd_t **dst, void **cbuf_handle, uint64_t *c_len,
     uint8_t level, boolean_t *local_offload)
 {
 	size_t d_len;
@@ -128,6 +128,12 @@ zia_compress_impl(const dpusm_uf_t *dpusm, zia_props_t *props,
 	/* nothing to offload, so just allocate space */
 	*cbuf_handle = zia_alloc(props->provider,
 	    s_len);
+	/*
+	 * Don't care about return value because if handle
+	 * association wasn't completed, it'll still work normally
+	 */
+	dpusm->associate_handle(*cbuf_handle, ABD_LINEAR_BUF(*dst));
+
 	if (!*cbuf_handle) {
 		return (ZIA_ERROR);
 	}

module/zia-software-provider/kernel_offloader.c

@@ -363,6 +363,24 @@ kernel_offloader_free(void *handle)
 	return (DPUSM_OK);
 }
 
+int
+kernel_offloader_associate_handle(void *handle, void *ptr) {
+	koh_t *koh = (koh_t *)unswizzle(handle);
+	if (!koh) {
+		return (KERNEL_OFFLOADER_ERROR);
+	}
+
+	write_lock(&rwlock);
+	active_size -= koh->size;
+	active_actual -= koh->size;
+	kfree(koh->ptr);
+	write_unlock(&rwlock);
+
+	koh->ptr = ptr;
+	koh->type = KOH_REFERENCE;
+	return (KERNEL_OFFLOADER_OK);
+}
+
 int
 kernel_offloader_copy_from_generic(void *handle, size_t offset,
     const void *src, size_t size)

module/zia-software-provider/kernel_offloader.h

@@ -106,6 +106,7 @@ void *kernel_offloader_alloc(size_t size);
 void *kernel_offloader_alloc_ref(void *src, size_t offset, size_t size);
 int kernel_offloader_get_size(void *handle, size_t *size, size_t *actual);
 int kernel_offloader_free(void *handle);
+int kernel_offloader_associate_handle(void *handle, void *ptr);
 int kernel_offloader_copy_from_generic(void *handle, size_t offset,
     const void *src, size_t size);
 int kernel_offloader_copy_from_scatterlist(void *ptr, size_t offset,

module/zia-software-provider/provider.c

@@ -196,6 +196,12 @@ sw_provider_get_size(void *handle, size_t *size, size_t *actual)
 	    size, actual)));
 }
 
+static int
+sw_provider_associate_handle(void *handle, void *ptr) {
+	return (translate_rc(kernel_offloader_associate_handle(handle,
+	    ptr)));
+}
+
 static int
 sw_provider_copy_from_generic(dpusm_mv_t *mv, const void *buf, size_t size)
 {
@@ -409,6 +415,7 @@ static const dpusm_pf_t sw_provider_functions = {
 	.alloc_ref            = kernel_offloader_alloc_ref,
 	.get_size             = sw_provider_get_size,
 	.free                 = kernel_offloader_free,
+	.associate_handle     = sw_provider_associate_handle,
 	.copy                 = {
 	                            .from = {
 	                                        .generic      = sw_provider_copy_from_generic,