Skip to content

Commit f738601

Browse files
kszafrankszafran
committed
Parallelise mock generation
1 parent fa81846 commit f738601

File tree

3 files changed

+129
-87
lines changed

3 files changed

+129
-87
lines changed

‎internal/cmd/mockery.go‎

Lines changed: 78 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"fmt"
77
"os"
88
"strings"
9+
"sync"
910

1011
"github.com/vektra/mockery/v3/config"
1112
"github.com/vektra/mockery/v3/internal"
@@ -161,7 +162,7 @@ func (i *InterfaceCollection) Append(ctx context.Context, iface *internal.Interf
161162
}
162163

163164
func (r *RootApp) Run() error {
164-
remoteTemplateCache := make(map[string]*internal.RemoteTemplate)
165+
remoteTemplateCache := pkg.NewRemoteTemplateCache()
165166

166167
log, err := logging.GetLogger(*r.Config.LogLevel)
167168
if err != nil {
@@ -296,66 +297,20 @@ func (r *RootApp) Run() error {
296297
}
297298
}
298299

300+
var wg sync.WaitGroup
299301
for outFilePath, interfacesInFile := range mockFileToInterfaces {
300-
fileLog := log.With().Str("file", outFilePath).Logger()
301-
fileCtx := fileLog.WithContext(ctx)
302+
wg.Add(1)
303+
go func() {
304+
defer wg.Done()
302305

303-
fileLog.Debug().Int("interfaces-in-file-len", len(interfacesInFile.interfaces)).Msgf("%v", interfacesInFile)
304-
305-
packageConfig, err := r.Config.GetPackageConfig(fileCtx, interfacesInFile.srcPkgPath)
306-
if err != nil {
307-
return err
308-
}
309-
if err := packageConfig.Config.ParseTemplates(ctx, "", "", interfacesInFile.srcPkg); err != nil {
310-
return err
311-
}
312-
313-
generator, err := pkg.NewTemplateGenerator(
314-
fileCtx,
315-
interfacesInFile.srcPkg,
316-
interfacesInFile.outFilePath.Parent(),
317-
*packageConfig.Config.Template,
318-
*packageConfig.Config.TemplateSchema,
319-
*packageConfig.Config.RequireTemplateSchemaExists,
320-
remoteTemplateCache,
321-
pkg.Formatter(*r.Config.Formatter),
322-
packageConfig.Config,
323-
interfacesInFile.outPkgName,
324-
)
325-
if err != nil {
326-
return err
327-
}
328-
fileLog.Info().Msg("Executing template")
329-
templateBytes, err := generator.Generate(
330-
fileCtx,
331-
interfacesInFile.interfaces,
332-
interfacesInFile.srcPkg.Name,
333-
interfacesInFile.srcPkg.PkgPath,
334-
)
335-
if err != nil {
336-
return err
337-
}
338-
339-
outFile := pathlib.NewPath(outFilePath)
340-
if err := outFile.Parent().MkdirAll(); err != nil {
341-
log.Err(err).Msg("failed to mkdir parent directories of mock file")
342-
return stackerr.NewStackErr(err)
343-
}
344-
fileLog.Info().Msg("Writing template to file")
345-
outFileExists, err := outFile.Exists()
346-
if err != nil {
347-
fileLog.Err(err).Msg("can't determine if outfile exists")
348-
return fmt.Errorf("determining if outfile exists: %w", err)
349-
}
350-
if outFileExists && !*packageConfig.Config.ForceFileWrite {
351-
fileLog.Error().Bool("force-file-write", *packageConfig.Config.ForceFileWrite).Msg("output file exists, can't write mocks")
352-
return fmt.Errorf("outfile exists")
353-
}
354-
355-
if err := outFile.WriteFile(templateBytes); err != nil {
356-
return stackerr.NewStackErr(err)
357-
}
306+
fileLog := log.With().Str("file", outFilePath).Logger()
307+
err := r.generate(log.WithContext(ctx), fileLog, outFilePath, interfacesInFile, remoteTemplateCache)
308+
if err != nil {
309+
fileLog.Err(err).Msg("error generating mocks")
310+
}
311+
}()
358312
}
313+
wg.Wait()
359314

360315
// The loop above could exit early, so sometimes warnings won't be shown
361316
// until other errors are fixed
@@ -375,3 +330,68 @@ func (r *RootApp) Run() error {
375330

376331
return nil
377332
}
333+
334+
func (r *RootApp) generate(
335+
ctx context.Context,
336+
log zerolog.Logger,
337+
outFilePath string,
338+
interfacesInFile *InterfaceCollection,
339+
remoteTemplateCache *pkg.RemoteTemplateCache,
340+
) error {
341+
log.Debug().Int("interfaces-in-file-len", len(interfacesInFile.interfaces)).Msgf("%v", interfacesInFile)
342+
343+
packageConfig, err := r.Config.GetPackageConfig(ctx, interfacesInFile.srcPkgPath)
344+
if err != nil {
345+
return err
346+
}
347+
if err := packageConfig.Config.ParseTemplates(ctx, "", "", interfacesInFile.srcPkg); err != nil {
348+
return err
349+
}
350+
351+
generator, err := pkg.NewTemplateGenerator(
352+
ctx,
353+
interfacesInFile.srcPkg,
354+
interfacesInFile.outFilePath.Parent(),
355+
*packageConfig.Config.Template,
356+
*packageConfig.Config.TemplateSchema,
357+
*packageConfig.Config.RequireTemplateSchemaExists,
358+
remoteTemplateCache,
359+
pkg.Formatter(*r.Config.Formatter),
360+
packageConfig.Config,
361+
interfacesInFile.outPkgName,
362+
)
363+
if err != nil {
364+
return err
365+
}
366+
log.Info().Msg("Executing template")
367+
templateBytes, err := generator.Generate(
368+
ctx,
369+
interfacesInFile.interfaces,
370+
interfacesInFile.srcPkg.Name,
371+
interfacesInFile.srcPkg.PkgPath,
372+
)
373+
if err != nil {
374+
return err
375+
}
376+
377+
outFile := pathlib.NewPath(outFilePath)
378+
if err := outFile.Parent().MkdirAll(); err != nil {
379+
log.Err(err).Msg("failed to mkdir parent directories of mock file")
380+
return stackerr.NewStackErr(err)
381+
}
382+
log.Info().Msg("Writing template to file")
383+
outFileExists, err := outFile.Exists()
384+
if err != nil {
385+
log.Err(err).Msg("can't determine if outfile exists")
386+
return fmt.Errorf("determining if outfile exists: %w", err)
387+
}
388+
if outFileExists && !*packageConfig.Config.ForceFileWrite {
389+
log.Error().Bool("force-file-write", *packageConfig.Config.ForceFileWrite).Msg("output file exists, can't write mocks")
390+
return fmt.Errorf("outfile exists")
391+
}
392+
393+
if err := outFile.WriteFile(templateBytes); err != nil {
394+
return stackerr.NewStackErr(err)
395+
}
396+
return nil
397+
}

‎internal/remote_template.go‎

Lines changed: 48 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"io"
77
"net/http"
88
"strings"
9+
"sync"
910

1011
"github.com/chigopher/pathlib"
1112
"github.com/rs/zerolog"
@@ -53,19 +54,45 @@ func download(ctx context.Context, url string) (string, error) {
5354
return "", fmt.Errorf("unsupported protocol specifier in %s", url)
5455
}
5556

57+
type RemoteTemplateCache struct {
58+
cache map[string]*RemoteTemplate
59+
mu sync.Mutex
60+
}
61+
62+
func NewRemoteTemplateCache() *RemoteTemplateCache {
63+
return &RemoteTemplateCache{
64+
cache: make(map[string]*RemoteTemplate),
65+
}
66+
}
67+
68+
func (c *RemoteTemplateCache) Get(templateURL, schemaURL string) *RemoteTemplate {
69+
c.mu.Lock()
70+
defer c.mu.Unlock()
71+
72+
if cached, ok := c.cache[templateURL]; !ok {
73+
template := NewRemoteTemplate(templateURL, schemaURL)
74+
c.cache[templateURL] = template
75+
return template
76+
} else {
77+
return cached
78+
}
79+
}
80+
5681
type RemoteTemplate struct {
57-
templateURL string
58-
templateString string
59-
templateDownloaded bool
82+
templateURL string
83+
templateString string
84+
templateOnce sync.Once
85+
templateErr error
6086

61-
schemaURL string
62-
schema *gojsonschema.Schema
63-
schemaDownloaded bool
87+
schemaURL string
88+
schema *gojsonschema.Schema
89+
schemaOnce sync.Once
90+
schemaErr error
6491

6592
requireSchemaExists bool
6693
}
6794

68-
func NewRemoteTemplate(templateURL string, schemaURL string) *RemoteTemplate {
95+
func NewRemoteTemplate(templateURL, schemaURL string) *RemoteTemplate {
6996
return &RemoteTemplate{
7097
templateURL: templateURL,
7198
schemaURL: schemaURL,
@@ -75,15 +102,15 @@ func NewRemoteTemplate(templateURL string, schemaURL string) *RemoteTemplate {
75102
// Template will return the template string. It downloads the remote template once
76103
// and caches the result for future calls.
77104
func (r *RemoteTemplate) Template(ctx context.Context) (string, error) {
78-
var err error
79-
if !r.templateDownloaded {
80-
r.templateDownloaded = true
105+
r.templateOnce.Do(func() {
106+
var err error
81107
r.templateString, err = download(ctx, r.templateURL)
82108
if err != nil {
83-
return "", fmt.Errorf("downloading template: %w", err)
109+
r.templateErr = fmt.Errorf("downloading template: %w", err)
84110
}
85-
}
86-
return r.templateString, nil
111+
})
112+
113+
return r.templateString, r.templateErr
87114
}
88115

89116
// Schema returns the JSON Schema as a string. It downloads the remote schema once
@@ -93,18 +120,20 @@ func (r *RemoteTemplate) Schema(ctx context.Context) (*gojsonschema.Schema, erro
93120
log.UpdateContext(func(c zerolog.Context) zerolog.Context {
94121
return c.Str("remote-template", r.templateURL)
95122
})
96-
if !r.schemaDownloaded {
123+
124+
r.schemaOnce.Do(func() {
97125
log.Debug().Msg("schema not downloaded before")
98-
r.schemaDownloaded = true
99126
schemaString, err := download(ctx, r.schemaURL)
100127
if err != nil {
101128
log.Debug().Err(err).Msg("schema download encountered error")
102-
return nil, fmt.Errorf("downloading schema: %w", err)
129+
r.schemaErr = fmt.Errorf("downloading schema: %w", err)
130+
return
103131
}
104132
r.schema, err = gojsonschema.NewSchema(gojsonschema.NewStringLoader(schemaString))
105133
if err != nil {
106-
return nil, fmt.Errorf("creating JSON schema: %w", err)
134+
r.schemaErr = fmt.Errorf("creating JSON schema: %w", err)
107135
}
108-
}
109-
return r.schema, nil
136+
})
137+
138+
return r.schema, r.schemaErr
110139
}

‎internal/template_generator.go‎

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,7 @@ type TemplateGenerator struct {
126126
templateSchema string
127127
pkgConfig *config.Config
128128
pkgName string
129-
remoteTemplateCache map[string]*RemoteTemplate
129+
remoteTemplateCache *RemoteTemplateCache
130130
}
131131

132132
func NewTemplateGenerator(
@@ -136,7 +136,7 @@ func NewTemplateGenerator(
136136
templateName string,
137137
templateSchema string,
138138
requireSchemaExists bool,
139-
remoteTemplateCache map[string]*RemoteTemplate,
139+
remoteTemplateCache *RemoteTemplateCache,
140140
formatter Formatter,
141141
pkgConfig *config.Config,
142142
pkgName string,
@@ -344,14 +344,7 @@ func (g *TemplateGenerator) getTemplate(ctx context.Context) (string, *gojsonsch
344344
if !strings.HasPrefix(g.templateName, protocol) {
345345
continue
346346
}
347-
var remoteTemplate *RemoteTemplate
348-
if cachedRemoteTemplate, ok := g.remoteTemplateCache[g.templateName]; !ok {
349-
remoteTemplate = NewRemoteTemplate(g.templateName, g.templateSchema)
350-
g.remoteTemplateCache[g.templateName] = remoteTemplate
351-
} else {
352-
remoteTemplate = cachedRemoteTemplate
353-
}
354-
347+
remoteTemplate := g.remoteTemplateCache.Get(g.templateName, g.templateSchema)
355348
templateString, err := remoteTemplate.Template(ctx)
356349
if err != nil {
357350
log.Error().Msg("could not download template")

0 commit comments

Comments
 (0)