hugo/lazy/init.go
Bjørn Erik Pedersen 3bbeb5688c Fix "context canceled" with partial
Make sure the context used for timeouts isn't created based on the incoming
context, as we have cases where this can cancel the context prematurely.

Fixes #10789
2023-03-04 21:29:05 +01:00

209 lines
4.7 KiB
Go

// Copyright 2019 The Hugo Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package lazy
import (
"context"
"sync"
"sync/atomic"
"time"
"errors"
)
// New creates a new empty Init.
func New() *Init {
return &Init{}
}
// Init holds a graph of lazily initialized dependencies.
type Init struct {
// Used mainly for testing.
initCount uint64
mu sync.Mutex
prev *Init
children []*Init
init onceMore
out any
err error
f func(context.Context) (any, error)
}
// Add adds a func as a new child dependency.
func (ini *Init) Add(initFn func(context.Context) (any, error)) *Init {
if ini == nil {
ini = New()
}
return ini.add(false, initFn)
}
// InitCount gets the number of this this Init has been initialized.
func (ini *Init) InitCount() int {
i := atomic.LoadUint64(&ini.initCount)
return int(i)
}
// AddWithTimeout is same as Add, but with a timeout that aborts initialization.
func (ini *Init) AddWithTimeout(timeout time.Duration, f func(ctx context.Context) (any, error)) *Init {
return ini.Add(func(ctx context.Context) (any, error) {
return ini.withTimeout(ctx, timeout, f)
})
}
// Branch creates a new dependency branch based on an existing and adds
// the given dependency as a child.
func (ini *Init) Branch(initFn func(context.Context) (any, error)) *Init {
if ini == nil {
ini = New()
}
return ini.add(true, initFn)
}
// BranchdWithTimeout is same as Branch, but with a timeout.
func (ini *Init) BranchWithTimeout(timeout time.Duration, f func(ctx context.Context) (any, error)) *Init {
return ini.Branch(func(ctx context.Context) (any, error) {
return ini.withTimeout(ctx, timeout, f)
})
}
// Do initializes the entire dependency graph.
func (ini *Init) Do(ctx context.Context) (any, error) {
if ini == nil {
panic("init is nil")
}
ini.init.Do(func() {
atomic.AddUint64(&ini.initCount, 1)
prev := ini.prev
if prev != nil {
// A branch. Initialize the ancestors.
if prev.shouldInitialize() {
_, err := prev.Do(ctx)
if err != nil {
ini.err = err
return
}
} else if prev.inProgress() {
// Concurrent initialization. The following init func
// may depend on earlier state, so wait.
prev.wait()
}
}
if ini.f != nil {
ini.out, ini.err = ini.f(ctx)
}
for _, child := range ini.children {
if child.shouldInitialize() {
_, err := child.Do(ctx)
if err != nil {
ini.err = err
return
}
}
}
})
ini.wait()
return ini.out, ini.err
}
// TODO(bep) investigate if we can use sync.Cond for this.
func (ini *Init) wait() {
var counter time.Duration
for !ini.init.Done() {
counter += 10
if counter > 600000000 {
panic("BUG: timed out in lazy init")
}
time.Sleep(counter * time.Microsecond)
}
}
func (ini *Init) inProgress() bool {
return ini != nil && ini.init.InProgress()
}
func (ini *Init) shouldInitialize() bool {
return !(ini == nil || ini.init.Done() || ini.init.InProgress())
}
// Reset resets the current and all its dependencies.
func (ini *Init) Reset() {
mu := ini.init.ResetWithLock()
ini.err = nil
defer mu.Unlock()
for _, d := range ini.children {
d.Reset()
}
}
func (ini *Init) add(branch bool, initFn func(context.Context) (any, error)) *Init {
ini.mu.Lock()
defer ini.mu.Unlock()
if branch {
return &Init{
f: initFn,
prev: ini,
}
}
ini.checkDone()
ini.children = append(ini.children, &Init{
f: initFn,
})
return ini
}
func (ini *Init) checkDone() {
if ini.init.Done() {
panic("init cannot be added to after it has run")
}
}
func (ini *Init) withTimeout(ctx context.Context, timeout time.Duration, f func(ctx context.Context) (any, error)) (any, error) {
// Create a new context with a timeout not connected to the incoming context.
waitCtx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
c := make(chan verr, 1)
go func() {
v, err := f(ctx)
select {
case <-waitCtx.Done():
return
default:
c <- verr{v: v, err: err}
}
}()
select {
case <-waitCtx.Done():
return nil, errors.New("timed out initializing value. You may have a circular loop in a shortcode, or your site may have resources that take longer to build than the `timeout` limit in your Hugo config file.")
case ve := <-c:
return ve.v, ve.err
}
}
type verr struct {
v any
err error
}