luciano/d2
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

don't use xmain.State in imgbundler library

Browse source
This commit is contained in:
Gavin Nishizawa 2023-10-19 16:06:11 -07:00
parent 0e8fa7cd8f
commit eba687fbbe
No known key found for this signature in database
GPG key ID: AE3B177777CE55CD
3 changed files with 39 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
d2cli/main.go
View file

@ -748,10 +748,11 @@ func _render(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, opts
return svg, err return svg, err
} }
svg, bundleErr := imgbundler.BundleLocal(ctx, ms, svg) cacheImages := ms.Env.Getenv("IMG_CACHE") == "1"
svg, bundleErr := imgbundler.BundleLocal(ctx, svg, cacheImages)
if bundle { if bundle {
var bundleErr2 error var bundleErr2 error
svg, bundleErr2 = imgbundler.BundleRemote(ctx, ms, svg) svg, bundleErr2 = imgbundler.BundleRemote(ctx, svg, cacheImages)
bundleErr = multierr.Combine(bundleErr, bundleErr2) bundleErr = multierr.Combine(bundleErr, bundleErr2)
} }
if forceAppendix && !toPNG { if forceAppendix && !toPNG {
@ -764,7 +765,7 @@ func _render(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, opts
if !bundle { if !bundle {
var bundleErr2 error var bundleErr2 error
svg, bundleErr2 = imgbundler.BundleRemote(ctx, ms, svg) svg, bundleErr2 = imgbundler.BundleRemote(ctx, svg, cacheImages)
bundleErr = multierr.Combine(bundleErr, bundleErr2) bundleErr = multierr.Combine(bundleErr, bundleErr2)
} }
@ -833,8 +834,9 @@ func renderPDF(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, opt
return svg, err return svg, err
} }
svg, bundleErr := imgbundler.BundleLocal(ctx, ms, svg) cacheImages := ms.Env.Getenv("IMG_CACHE") == "1"
svg, bundleErr2 := imgbundler.BundleRemote(ctx, ms, svg) svg, bundleErr := imgbundler.BundleLocal(ctx, svg, cacheImages)
svg, bundleErr2 := imgbundler.BundleRemote(ctx, svg, cacheImages)
bundleErr = multierr.Combine(bundleErr, bundleErr2) bundleErr = multierr.Combine(bundleErr, bundleErr2)
if bundleErr != nil { if bundleErr != nil {
return svg, bundleErr return svg, bundleErr
@ -933,8 +935,9 @@ func renderPPTX(ctx context.Context, ms *xmain.State, presentation *pptx.Present
return nil, err return nil, err
} }
svg, bundleErr := imgbundler.BundleLocal(ctx, ms, svg) cacheImages := ms.Env.Getenv("IMG_CACHE") == "1"
svg, bundleErr2 := imgbundler.BundleRemote(ctx, ms, svg) svg, bundleErr := imgbundler.BundleLocal(ctx, svg, cacheImages)
svg, bundleErr2 := imgbundler.BundleRemote(ctx, svg, cacheImages)
bundleErr = multierr.Combine(bundleErr, bundleErr2) bundleErr = multierr.Combine(bundleErr, bundleErr2)
if bundleErr != nil { if bundleErr != nil {
return nil, bundleErr return nil, bundleErr
@ -1178,8 +1181,9 @@ func renderPNGsForGIF(ctx context.Context, ms *xmain.State, plugin d2plugin.Plug
return nil, nil, err return nil, nil, err
} }
svg, bundleErr := imgbundler.BundleLocal(ctx, ms, svg) cacheImages := ms.Env.Getenv("IMG_CACHE") == "1"
svg, bundleErr2 := imgbundler.BundleRemote(ctx, ms, svg) svg, bundleErr := imgbundler.BundleLocal(ctx, svg, cacheImages)
svg, bundleErr2 := imgbundler.BundleRemote(ctx, svg, cacheImages)
bundleErr = multierr.Combine(bundleErr, bundleErr2) bundleErr = multierr.Combine(bundleErr, bundleErr2)
if bundleErr != nil { if bundleErr != nil {
return nil, nil, bundleErr return nil, nil, bundleErr

32
lib/imgbundler/imgbundler.go
View file

@ -19,8 +19,8 @@ import (
"golang.org/x/xerrors" "golang.org/x/xerrors"
"oss.terrastruct.com/d2/lib/log"
"oss.terrastruct.com/util-go/xdefer" "oss.terrastruct.com/util-go/xdefer"
"oss.terrastruct.com/util-go/xmain"
) )
var imgCache sync.Map var imgCache sync.Map
@ -29,12 +29,12 @@ const maxImageSize int64 = 1 << 25 // 33_554_432
var imageRegex = regexp.MustCompile(`<image href="([^"]+)"`) var imageRegex = regexp.MustCompile(`<image href="([^"]+)"`)
func BundleLocal(ctx context.Context, ms *xmain.State, in []byte) ([]byte, error) { func BundleLocal(ctx context.Context, in []byte, cacheImages bool) ([]byte, error) {
return bundle(ctx, ms, in, false) return bundle(ctx, in, false, cacheImages)
} }
func BundleRemote(ctx context.Context, ms *xmain.State, in []byte) ([]byte, error) { func BundleRemote(ctx context.Context, in []byte, cacheImages bool) ([]byte, error) {
return bundle(ctx, ms, in, true) return bundle(ctx, in, true, cacheImages)
} }
type repl struct { type repl struct {
@ -42,7 +42,7 @@ type repl struct {
to []byte to []byte
} }
func bundle(ctx context.Context, ms *xmain.State, svg []byte, isRemote bool) (_ []byte, err error) { func bundle(ctx context.Context, svg []byte, isRemote, cacheImages bool) (_ []byte, err error) {
if isRemote { if isRemote {
defer xdefer.Errorf(&err, "failed to bundle remote images") defer xdefer.Errorf(&err, "failed to bundle remote images")
} else { } else {
@ -54,7 +54,7 @@ func bundle(ctx context.Context, ms *xmain.State, svg []byte, isRemote bool) (_
ctx, cancel := context.WithTimeout(ctx, time.Minute*5) ctx, cancel := context.WithTimeout(ctx, time.Minute*5)
defer cancel() defer cancel()
return runWorkers(ctx, ms, svg, imgs, isRemote) return runWorkers(ctx, svg, imgs, isRemote, cacheImages)
} }
// filterImageElements finds all unique image elements in imgs that are // filterImageElements finds all unique image elements in imgs that are
@ -84,7 +84,7 @@ func filterImageElements(imgs [][][]byte, isRemote bool) [][][]byte {
return imgs2 return imgs2
} }
func runWorkers(ctx context.Context, ms *xmain.State, svg []byte, imgs [][][]byte, isRemote bool) (_ []byte, err error) { func runWorkers(ctx context.Context, svg []byte, imgs [][][]byte, isRemote, cacheImages bool) (_ []byte, err error) {
var wg sync.WaitGroup var wg sync.WaitGroup
replc := make(chan repl) replc := make(chan repl)
@ -111,9 +111,9 @@ func runWorkers(ctx context.Context, ms *xmain.State, svg []byte, imgs [][][]byt
<-sema <-sema
}() }()
bundledImage, err := worker(ctx, ms, img[1], isRemote) bundledImage, err := worker(ctx, img[1], isRemote, cacheImages)
if err != nil { if err != nil {
ms.Log.Error.Printf("failed to bundle %s: %v", img[1], err) log.Error(ctx, fmt.Sprintf("failed to bundle %s: %v", img[1], err))
errhrefsMu.Lock() errhrefsMu.Lock()
errhrefs = append(errhrefs, string(img[1])) errhrefs = append(errhrefs, string(img[1]))
errhrefsMu.Unlock() errhrefsMu.Unlock()
@ -137,7 +137,7 @@ func runWorkers(ctx context.Context, ms *xmain.State, svg []byte, imgs [][][]byt
case <-ctx.Done(): case <-ctx.Done():
return svg, xerrors.Errorf("failed to wait for workers: %w", ctx.Err()) return svg, xerrors.Errorf("failed to wait for workers: %w", ctx.Err())
case <-t.C: case <-t.C:
ms.Log.Info.Printf("fetching images...") log.Info(ctx, "fetching images...")
case repl, ok := <-replc: case repl, ok := <-replc:
if !ok { if !ok {
if len(errhrefs) > 0 { if len(errhrefs) > 0 {
@ -150,8 +150,8 @@ func runWorkers(ctx context.Context, ms *xmain.State, svg []byte, imgs [][][]byt
} }
} }
func worker(ctx context.Context, ms *xmain.State, href []byte, isRemote bool) ([]byte, error) { func worker(ctx context.Context, href []byte, isRemote, cacheImages bool) ([]byte, error) {
if ms.Env.Getenv("IMG_CACHE") == "1" { if cacheImages {
if hit, ok := imgCache.Load(string(href)); ok { if hit, ok := imgCache.Load(string(href)); ok {
return hit.([]byte), nil return hit.([]byte), nil
} }
@ -160,10 +160,10 @@ func worker(ctx context.Context, ms *xmain.State, href []byte, isRemote bool) ([
var mimeType string var mimeType string
var err error var err error
if isRemote { if isRemote {
ms.Log.Debug.Printf("fetching %s remotely", string(href)) log.Debug(ctx, fmt.Sprintf("fetching %s remotely", string(href)))
buf, mimeType, err = httpGet(ctx, html.UnescapeString(string(href))) buf, mimeType, err = httpGet(ctx, html.UnescapeString(string(href)))
} else { } else {
ms.Log.Debug.Printf("reading %s from disk", string(href)) log.Debug(ctx, fmt.Sprintf("reading %s from disk", string(href)))
buf, err = os.ReadFile(html.UnescapeString(string(href))) buf, err = os.ReadFile(html.UnescapeString(string(href)))
} }
if err != nil { if err != nil {
@ -177,7 +177,7 @@ func worker(ctx context.Context, ms *xmain.State, href []byte, isRemote bool) ([
b64 := base64.StdEncoding.EncodeToString(buf) b64 := base64.StdEncoding.EncodeToString(buf)
out := []byte(fmt.Sprintf(`<image href="data:%s;base64,%s"`, mimeType, b64)) out := []byte(fmt.Sprintf(`<image href="data:%s;base64,%s"`, mimeType, b64))
if ms.Env.Getenv("IMG_CACHE") == "1" { if cacheImages {
imgCache.Store(string(href), out) imgCache.Store(string(href), out)
} }
return out, nil return out, nil

22
lib/imgbundler/imgbundler_test.go
View file

@ -109,7 +109,7 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
return respRecorder.Result() return respRecorder.Result()
}) })
out, err := BundleRemote(ctx, ms, []byte(sampleSVG)) out, err := BundleRemote(ctx, []byte(sampleSVG), false)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@ -133,7 +133,7 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
respRecorder.WriteHeader(200) respRecorder.WriteHeader(200)
return respRecorder.Result() return respRecorder.Result()
}) })
_, err = BundleRemote(ctx, ms, []byte(sampleSVG)) _, err = BundleRemote(ctx, []byte(sampleSVG), false)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@ -148,7 +148,7 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
respRecorder.WriteHeader(200) respRecorder.WriteHeader(200)
return respRecorder.Result() return respRecorder.Result()
}) })
_, err = BundleRemote(ctx, ms, []byte(sampleSVG)) _, err = BundleRemote(ctx, []byte(sampleSVG), false)
if err == nil { if err == nil {
t.Fatal("expected error") t.Fatal("expected error")
} }
@ -160,7 +160,7 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
respRecorder.WriteHeader(500) respRecorder.WriteHeader(500)
return respRecorder.Result() return respRecorder.Result()
}) })
_, err = BundleRemote(ctx, ms, []byte(sampleSVG)) _, err = BundleRemote(ctx, []byte(sampleSVG), false)
if err == nil { if err == nil {
t.Fatal("expected error") t.Fatal("expected error")
} }
@ -215,7 +215,7 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
Env: xos.NewEnv(os.Environ()), Env: xos.NewEnv(os.Environ()),
} }
ms.Log = cmdlog.NewTB(ms.Env, t) ms.Log = cmdlog.NewTB(ms.Env, t)
out, err := BundleLocal(ctx, ms, []byte(sampleSVG)) out, err := BundleLocal(ctx, []byte(sampleSVG), false)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@ -285,7 +285,7 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
return respRecorder.Result() return respRecorder.Result()
}) })
out, err := BundleRemote(ctx, ms, []byte(sampleSVG)) out, err := BundleRemote(ctx, []byte(sampleSVG), false)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@ -351,25 +351,23 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
}) })
// Using a cache, imgs are not refetched on multiple runs // Using a cache, imgs are not refetched on multiple runs
ms.Env.Setenv("IMG_CACHE", "1") _, err := BundleRemote(ctx, []byte(sampleSVG), true)
_, err := BundleRemote(ctx, ms, []byte(sampleSVG))
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
_, err = BundleRemote(ctx, ms, []byte(sampleSVG)) _, err = BundleRemote(ctx, []byte(sampleSVG), true)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
tassert.Equal(t, 1, count) tassert.Equal(t, 1, count)
// With cache disabled, it refetches // With cache disabled, it refetches
ms.Env.Setenv("IMG_CACHE", "0")
count = 0 count = 0
_, err = BundleRemote(ctx, ms, []byte(sampleSVG)) _, err = BundleRemote(ctx, []byte(sampleSVG), false)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
_, err = BundleRemote(ctx, ms, []byte(sampleSVG)) _, err = BundleRemote(ctx, []byte(sampleSVG), false)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }