perf: switch to sync.Map-based pointer store

Author: CoolSpring8

builder.go

@@ -13,23 +13,25 @@ extern lol_html_rewriter_directive_t callback_doc_end(lol_html_doc_end_t *doc_en
 import "C"
 import (
 	"unsafe"
-
-	"github.com/mattn/go-pointer"
 )
 
 // rewriterBuilder is used to build a rewriter.
 type rewriterBuilder struct {
 	rb       *C.lol_html_rewriter_builder_t
 	pointers []unsafe.Pointer
+	built    bool // this builder has built at least one writer
 }
 
 func newRewriterBuilder() *rewriterBuilder {
-	return &rewriterBuilder{rb: C.lol_html_rewriter_builder_new(), pointers: nil}
+	return &rewriterBuilder{rb: C.lol_html_rewriter_builder_new(), pointers: nil, built: false}
 }
 
 func (rb *rewriterBuilder) Free() {
 	if rb != nil {
 		C.lol_html_rewriter_builder_free(rb.rb)
+		if !rb.built {
+			unrefPointers(rb.pointers)
+		}
 	}
 }
 
@@ -52,10 +54,10 @@ func (rb *rewriterBuilder) AddDocumentContentHandlers(
 	if documentEndHandler != nil {
 		cCallbackDocumentEndPointer = (*[0]byte)(C.callback_doc_end)
 	}
-	doctypeHandlerPointer := pointer.Save(doctypeHandler)
-	commentHandlerPointer := pointer.Save(commentHandler)
-	textChunkHandlerPointer := pointer.Save(textChunkHandler)
-	documentEndHandlerPointer := pointer.Save(documentEndHandler)
+	doctypeHandlerPointer := savePointer(doctypeHandler)
+	commentHandlerPointer := savePointer(commentHandler)
+	textChunkHandlerPointer := savePointer(textChunkHandler)
+	documentEndHandlerPointer := savePointer(documentEndHandler)
 	C.lol_html_rewriter_builder_add_document_content_handlers(
 		rb.rb,
 		cCallbackDoctypePointer,
@@ -92,9 +94,9 @@ func (rb *rewriterBuilder) AddElementContentHandlers(
 	if textChunkHandler != nil {
 		cCallbackTextChunkPointer = (*[0]byte)(C.callback_text_chunk)
 	}
-	elementHandlerPointer := pointer.Save(elementHandler)
-	commentHandlerPointer := pointer.Save(commentHandler)
-	textChunkHandlerPointer := pointer.Save(textChunkHandler)
+	elementHandlerPointer := savePointer(elementHandler)
+	commentHandlerPointer := savePointer(commentHandler)
+	textChunkHandlerPointer := savePointer(textChunkHandler)
 	C.lol_html_rewriter_builder_add_element_content_handlers(
 		rb.rb,
 		(*C.lol_html_selector_t)(selector),
@@ -116,7 +118,7 @@ func (rb *rewriterBuilder) Build(sink OutputSink, config Config) (*rewriter, err
 		preallocated_parsing_buffer_size: C.size_t(config.Memory.PreallocatedParsingBufferSize),
 		max_allowed_memory_usage:         C.size_t(config.Memory.MaxAllowedMemoryUsage),
 	}
-	p := pointer.Save(sink)
+	p := savePointer(sink)
 	r := C.lol_html_rewriter_build(
 		rb.rb,
 		encodingC,
@@ -127,6 +129,7 @@ func (rb *rewriterBuilder) Build(sink OutputSink, config Config) (*rewriter, err
 		C.bool(config.Strict),
 	)
 	if r != nil {
+		rb.built = true
 		return &rewriter{rw: r, pointers: rb.pointers}, nil
 	}
 	return nil, getError()

config.go

@@ -5,7 +5,6 @@ package lolhtml
 */
 import "C"
 import (
-	"github.com/mattn/go-pointer"
 	"unsafe"
 )
 
@@ -63,36 +62,36 @@ type Handlers struct {
 //export callbackSink
 func callbackSink(chunk *C.char, chunkLen C.size_t, userData unsafe.Pointer) {
 	c := C.GoBytes(unsafe.Pointer(chunk), C.int(chunkLen))
-	cb := pointer.Restore(userData).(OutputSink)
+	cb := restorePointer(userData).(OutputSink)
 	cb(c)
 }
 
 //export callbackDoctype
 func callbackDoctype(doctype *Doctype, userData unsafe.Pointer) RewriterDirective {
-	cb := pointer.Restore(userData).(DoctypeHandlerFunc)
+	cb := restorePointer(userData).(DoctypeHandlerFunc)
 	return cb(doctype)
 }
 
 //export callbackComment
 func callbackComment(comment *Comment, userData unsafe.Pointer) RewriterDirective {
-	cb := pointer.Restore(userData).(CommentHandlerFunc)
+	cb := restorePointer(userData).(CommentHandlerFunc)
 	return cb(comment)
 }
 
 //export callbackTextChunk
 func callbackTextChunk(textChunk *TextChunk, userData unsafe.Pointer) RewriterDirective {
-	cb := pointer.Restore(userData).(TextChunkHandlerFunc)
+	cb := restorePointer(userData).(TextChunkHandlerFunc)
 	return cb(textChunk)
 }
 
 //export callbackElement
 func callbackElement(element *Element, userData unsafe.Pointer) RewriterDirective {
-	cb := pointer.Restore(userData).(ElementHandlerFunc)
+	cb := restorePointer(userData).(ElementHandlerFunc)
 	return cb(element)
 }
 
 //export callbackDocumentEnd
 func callbackDocumentEnd(documentEnd *DocumentEnd, userData unsafe.Pointer) RewriterDirective {
-	cb := pointer.Restore(userData).(DocumentEndHandlerFunc)
+	cb := restorePointer(userData).(DocumentEndHandlerFunc)
 	return cb(documentEnd)
 }

go.mod

@@ -1,5 +1,3 @@
 module github.com/coolspring8/go-lolhtml
 
 go 1.15
-
-require github.com/mattn/go-pointer v0.0.1

go.sum

@@ -1,2 +0,0 @@
-github.com/mattn/go-pointer v0.0.1 h1:n+XhsuGeVO6MEAp7xyEukFINEa+Quek5psIR/ylA6o0=
-github.com/mattn/go-pointer v0.0.1/go.mod h1:2zXcozF6qYGgmsG+SeTZz3oAbFLdD3OWqnUbNvJZAlc=

pointer.go

@@ -0,0 +1,58 @@
+package lolhtml
+
+// Credit to https://github.com/mattn/go-pointer.
+
+// #include <stdlib.h>
+import "C"
+import (
+	"sync"
+	"unsafe"
+)
+
+// sync.Map documentation states that it is optimized for "when the entry for a given key is only
+// ever written once but read many times, as in caches that only grow". My benchmarks show that sync.Map
+// version rewriter is slower in single-goroutine calls, but faster when used in multiple goroutines
+// (and personally I think the latter is more important).
+var store sync.Map
+
+func savePointer(v interface{}) unsafe.Pointer {
+	if v == nil {
+		return nil
+	}
+
+	ptr := C.malloc(C.size_t(1))
+	if ptr == nil {
+		panic(`can't allocate "cgo-pointer hack index pointer": ptr == nil`)
+	}
+
+	store.Store(ptr, v)
+
+	return ptr
+}
+
+func restorePointer(ptr unsafe.Pointer) (v interface{}) {
+	if ptr == nil {
+		return nil
+	}
+
+	if v, ok := store.Load(ptr); ok {
+		return v
+	}
+	return nil
+}
+
+func unrefPointer(ptr unsafe.Pointer) {
+	if ptr == nil {
+		return
+	}
+
+	store.Delete(ptr)
+
+	C.free(ptr)
+}
+
+func unrefPointers(ptrs []unsafe.Pointer) {
+	for _, ptr := range ptrs {
+		unrefPointer(ptr)
+	}
+}

rewriter.go

@@ -7,8 +7,6 @@ package lolhtml
 import "C"
 import (
 	"unsafe"
-
-	"github.com/mattn/go-pointer"
 )
 
 // rewriter represents an actual HTML rewriter.
@@ -55,8 +53,6 @@ func (r *rewriter) End() error {
 func (r *rewriter) Free() {
 	if r != nil {
 		C.lol_html_rewriter_free(r.rw)
-		for _, p := range r.pointers {
-			pointer.Unref(p)
-		}
+		unrefPointers(r.pointers)
 	}
 }