Skip to content

Commit

Permalink
Merge pull request #1678 from gavin-ts/update-imgbundler
Browse files Browse the repository at this point in the history 
remove imgbundler's xmain.State dependency
  • Loading branch information
@gavin-ts
gavin-ts authored Oct 20, 2023
2 parents 0e8fa7c + 32fbe7e commit 4cba5d2
Show file tree
Hide file tree
Showing 7 changed files with 153 additions and 115 deletions.
55 changes: 41 additions & 14 deletions d2cli/main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ import (
"oss.terrastruct.com/d2/lib/pdf"
"oss.terrastruct.com/d2/lib/png"
"oss.terrastruct.com/d2/lib/pptx"
"oss.terrastruct.com/d2/lib/simplelog"
"oss.terrastruct.com/d2/lib/textmeasure"
timelib "oss.terrastruct.com/d2/lib/time"
"oss.terrastruct.com/d2/lib/version"
Expand Down Expand Up @@ -434,7 +435,7 @@ func compile(ctx context.Context, ms *xmain.State, plugins []d2plugin.Plugin, la
if err != nil {
return nil, false, err
}
out, err := xgif.AnimatePNGs(ms, pngs, int(animateInterval))
out, err := AnimatePNGs(ms, pngs, int(animateInterval))
if err != nil {
return nil, false, err
}
Expand Down Expand Up @@ -748,10 +749,12 @@ func _render(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, opts
return svg, err
}

svg, bundleErr := imgbundler.BundleLocal(ctx, ms, svg)
cacheImages := ms.Env.Getenv("IMG_CACHE") == "1"
l := simplelog.FromCmdLog(ms.Log)
svg, bundleErr := imgbundler.BundleLocal(ctx, l, svg, cacheImages)
if bundle {
var bundleErr2 error
svg, bundleErr2 = imgbundler.BundleRemote(ctx, ms, svg)
svg, bundleErr2 = imgbundler.BundleRemote(ctx, l, svg, cacheImages)
bundleErr = multierr.Combine(bundleErr, bundleErr2)
}
if forceAppendix && !toPNG {
Expand All @@ -764,11 +767,11 @@ func _render(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, opts

if !bundle {
var bundleErr2 error
svg, bundleErr2 = imgbundler.BundleRemote(ctx, ms, svg)
svg, bundleErr2 = imgbundler.BundleRemote(ctx, l, svg, cacheImages)
bundleErr = multierr.Combine(bundleErr, bundleErr2)
}

out, err = png.ConvertSVG(ms, page, svg)
out, err = ConvertSVG(ms, page, svg)
if err != nil {
return svg, err
}
Expand Down Expand Up @@ -833,15 +836,17 @@ func renderPDF(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, opt
return svg, err
}

svg, bundleErr := imgbundler.BundleLocal(ctx, ms, svg)
svg, bundleErr2 := imgbundler.BundleRemote(ctx, ms, svg)
cacheImages := ms.Env.Getenv("IMG_CACHE") == "1"
l := simplelog.FromCmdLog(ms.Log)
svg, bundleErr := imgbundler.BundleLocal(ctx, l, svg, cacheImages)
svg, bundleErr2 := imgbundler.BundleRemote(ctx, l, svg, cacheImages)
bundleErr = multierr.Combine(bundleErr, bundleErr2)
if bundleErr != nil {
return svg, bundleErr
}
svg = appendix.Append(diagram, ruler, svg)

pngImg, err := png.ConvertSVG(ms, page, svg)
pngImg, err := ConvertSVG(ms, page, svg)
if err != nil {
return svg, err
}
Expand Down Expand Up @@ -933,16 +938,18 @@ func renderPPTX(ctx context.Context, ms *xmain.State, presentation *pptx.Present
return nil, err
}

svg, bundleErr := imgbundler.BundleLocal(ctx, ms, svg)
svg, bundleErr2 := imgbundler.BundleRemote(ctx, ms, svg)
cacheImages := ms.Env.Getenv("IMG_CACHE") == "1"
l := simplelog.FromCmdLog(ms.Log)
svg, bundleErr := imgbundler.BundleLocal(ctx, l, svg, cacheImages)
svg, bundleErr2 := imgbundler.BundleRemote(ctx, l, svg, cacheImages)
bundleErr = multierr.Combine(bundleErr, bundleErr2)
if bundleErr != nil {
return nil, bundleErr
}

svg = appendix.Append(diagram, ruler, svg)

pngImg, err := png.ConvertSVG(ms, page, svg)
pngImg, err := ConvertSVG(ms, page, svg)
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -1178,16 +1185,18 @@ func renderPNGsForGIF(ctx context.Context, ms *xmain.State, plugin d2plugin.Plug
return nil, nil, err
}

svg, bundleErr := imgbundler.BundleLocal(ctx, ms, svg)
svg, bundleErr2 := imgbundler.BundleRemote(ctx, ms, svg)
cacheImages := ms.Env.Getenv("IMG_CACHE") == "1"
l := simplelog.FromCmdLog(ms.Log)
svg, bundleErr := imgbundler.BundleLocal(ctx, l, svg, cacheImages)
svg, bundleErr2 := imgbundler.BundleRemote(ctx, l, svg, cacheImages)
bundleErr = multierr.Combine(bundleErr, bundleErr2)
if bundleErr != nil {
return nil, nil, bundleErr
}

svg = appendix.Append(diagram, ruler, svg)

pngImg, err := png.ConvertSVG(ms, page, svg)
pngImg, err := ConvertSVG(ms, page, svg)
if err != nil {
return nil, nil, err
}
Expand Down Expand Up @@ -1218,3 +1227,21 @@ func renderPNGsForGIF(ctx context.Context, ms *xmain.State, plugin d2plugin.Plug

return svg, pngs, nil
}

func ConvertSVG(ms *xmain.State, page playwright.Page, svg []byte) ([]byte, error) {
cancel := background.Repeat(func() {
ms.Log.Info.Printf("converting to PNG...")
}, time.Second*5)
defer cancel()

return png.ConvertSVG(page, svg)
}

func AnimatePNGs(ms *xmain.State, pngs [][]byte, animIntervalMs int) ([]byte, error) {
cancel := background.Repeat(func() {
ms.Log.Info.Printf("generating GIF...")
}, time.Second*5)
defer cancel()

return xgif.AnimatePNGs(pngs, animIntervalMs)
}
32 changes: 16 additions & 16 deletions lib/imgbundler/imgbundler.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,8 @@ import (

"golang.org/x/xerrors"

"oss.terrastruct.com/d2/lib/simplelog"
"oss.terrastruct.com/util-go/xdefer"
"oss.terrastruct.com/util-go/xmain"
)

var imgCache sync.Map
Expand All @@ -29,20 +29,20 @@ const maxImageSize int64 = 1 << 25 // 33_554_432

var imageRegex = regexp.MustCompile(`<image href="([^"]+)"`)

func BundleLocal(ctx context.Context, ms *xmain.State, in []byte) ([]byte, error) {
return bundle(ctx, ms, in, false)
func BundleLocal(ctx context.Context, l simplelog.Logger, in []byte, cacheImages bool) ([]byte, error) {
return bundle(ctx, l, in, false, cacheImages)
}

func BundleRemote(ctx context.Context, ms *xmain.State, in []byte) ([]byte, error) {
return bundle(ctx, ms, in, true)
func BundleRemote(ctx context.Context, l simplelog.Logger, in []byte, cacheImages bool) ([]byte, error) {
return bundle(ctx, l, in, true, cacheImages)
}

type repl struct {
from []byte
to []byte
}

func bundle(ctx context.Context, ms *xmain.State, svg []byte, isRemote bool) (_ []byte, err error) {
func bundle(ctx context.Context, l simplelog.Logger, svg []byte, isRemote, cacheImages bool) (_ []byte, err error) {
if isRemote {
defer xdefer.Errorf(&err, "failed to bundle remote images")
} else {
Expand All @@ -54,7 +54,7 @@ func bundle(ctx context.Context, ms *xmain.State, svg []byte, isRemote bool) (_
ctx, cancel := context.WithTimeout(ctx, time.Minute*5)
defer cancel()

return runWorkers(ctx, ms, svg, imgs, isRemote)
return runWorkers(ctx, l, svg, imgs, isRemote, cacheImages)
}

// filterImageElements finds all unique image elements in imgs that are
Expand Down Expand Up @@ -84,7 +84,7 @@ func filterImageElements(imgs [][][]byte, isRemote bool) [][][]byte {
return imgs2
}

func runWorkers(ctx context.Context, ms *xmain.State, svg []byte, imgs [][][]byte, isRemote bool) (_ []byte, err error) {
func runWorkers(ctx context.Context, l simplelog.Logger, svg []byte, imgs [][][]byte, isRemote, cacheImages bool) (_ []byte, err error) {
var wg sync.WaitGroup
replc := make(chan repl)

Expand All @@ -111,9 +111,9 @@ func runWorkers(ctx context.Context, ms *xmain.State, svg []byte, imgs [][][]byt
<-sema
}()

bundledImage, err := worker(ctx, ms, img[1], isRemote)
bundledImage, err := worker(ctx, l, img[1], isRemote, cacheImages)
if err != nil {
ms.Log.Error.Printf("failed to bundle %s: %v", img[1], err)
l.Error(fmt.Sprintf("failed to bundle %s: %v", img[1], err))
errhrefsMu.Lock()
errhrefs = append(errhrefs, string(img[1]))
errhrefsMu.Unlock()
Expand All @@ -137,7 +137,7 @@ func runWorkers(ctx context.Context, ms *xmain.State, svg []byte, imgs [][][]byt
case <-ctx.Done():
return svg, xerrors.Errorf("failed to wait for workers: %w", ctx.Err())
case <-t.C:
ms.Log.Info.Printf("fetching images...")
l.Info("fetching images...")
case repl, ok := <-replc:
if !ok {
if len(errhrefs) > 0 {
Expand All @@ -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) {
if ms.Env.Getenv("IMG_CACHE") == "1" {
func worker(ctx context.Context, l simplelog.Logger, href []byte, isRemote, cacheImages bool) ([]byte, error) {
if cacheImages {
if hit, ok := imgCache.Load(string(href)); ok {
return hit.([]byte), nil
}
Expand All @@ -160,10 +160,10 @@ func worker(ctx context.Context, ms *xmain.State, href []byte, isRemote bool) ([
var mimeType string
var err error
if isRemote {
ms.Log.Debug.Printf("fetching %s remotely", string(href))
l.Debug(fmt.Sprintf("fetching %s remotely", string(href)))
buf, mimeType, err = httpGet(ctx, html.UnescapeString(string(href)))
} else {
ms.Log.Debug.Printf("reading %s from disk", string(href))
l.Debug(fmt.Sprintf("reading %s from disk", string(href)))
buf, err = os.ReadFile(html.UnescapeString(string(href)))
}
if err != nil {
Expand All @@ -177,7 +177,7 @@ func worker(ctx context.Context, ms *xmain.State, href []byte, isRemote bool) ([
b64 := base64.StdEncoding.EncodeToString(buf)

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)
}
return out, nil
Expand Down
Loading

0 comments on commit 4cba5d2

Please sign in to comment.