narratio/internal/adapters/whisperx/multipart_body.go

package whisperx

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"path/filepath"
	"sync"
)

type multipartBodyFactory struct {
	ctx         context.Context
	audioPath   string
	openAudio   func(string) (io.ReadCloser, error)
	prefix      []byte
	suffix      []byte
	contentType string

	mu     sync.Mutex
	closed bool
	bodies []*multipartBody
}

func newMultipartBodyFactory(
	ctx context.Context,
	audioPath string,
	language string,
	openAudio func(string) (io.ReadCloser, error),
) (*multipartBodyFactory, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if openAudio == nil {
		return nil, fmt.Errorf("audio opener is required")
	}
	var framing bytes.Buffer
	writer := multipart.NewWriter(&framing)
	if _, err := writer.CreateFormFile("file", filepath.Base(audioPath)); err != nil {
		return nil, fmt.Errorf("create multipart file field: %w", err)
	}
	prefixLength := framing.Len()
	if err := writer.WriteField("language", language); err != nil {
		return nil, fmt.Errorf("write language form field: %w", err)
	}
	if err := writer.Close(); err != nil {
		return nil, fmt.Errorf("close multipart framing: %w", err)
	}
	framingBytes := framing.Bytes()
	return &multipartBodyFactory{
		ctx:         ctx,
		audioPath:   audioPath,
		openAudio:   openAudio,
		prefix:      append([]byte(nil), framingBytes[:prefixLength]...),
		suffix:      append([]byte(nil), framingBytes[prefixLength:]...),
		contentType: writer.FormDataContentType(),
	}, nil
}

func (factory *multipartBodyFactory) NewBody() (io.ReadCloser, error) {
	factory.mu.Lock()
	if factory.closed {
		factory.mu.Unlock()
		return nil, fmt.Errorf("multipart body factory is closed")
	}
	if err := factory.ctx.Err(); err != nil {
		factory.mu.Unlock()
		return nil, err
	}
	audio, err := factory.openAudio(factory.audioPath)
	if err != nil {
		factory.mu.Unlock()
		return nil, fmt.Errorf("open audio file %q: %w", factory.audioPath, err)
	}
	if err := factory.ctx.Err(); err != nil {
		factory.mu.Unlock()
		_ = audio.Close()
		return nil, err
	}
	source := &multipartAudioSource{audio: audio, closeDone: make(chan struct{})}
	body := &multipartBody{
		ctx:    factory.ctx,
		source: source,
		reader: io.MultiReader(bytes.NewReader(factory.prefix), source, bytes.NewReader(factory.suffix)),
	}
	factory.bodies = append(factory.bodies, body)
	factory.mu.Unlock()
	stopCancel := context.AfterFunc(factory.ctx, func() {
		_ = body.Close()
	})
	body.mu.Lock()
	body.stopCancel = stopCancel
	closed := body.closed
	body.mu.Unlock()
	if closed {
		stopCancel()
	}
	if err := factory.ctx.Err(); err != nil {
		_ = body.Close()
		return nil, err
	}
	return body, nil
}

func (factory *multipartBodyFactory) Close() error {
	factory.mu.Lock()
	if factory.closed {
		bodies := append([]*multipartBody(nil), factory.bodies...)
		factory.mu.Unlock()
		return multipartBodyCloseErrors(bodies)
	}
	factory.closed = true
	bodies := append([]*multipartBody(nil), factory.bodies...)
	factory.mu.Unlock()
	for _, body := range bodies {
		_ = body.Close()
	}
	return multipartBodyCloseErrors(bodies)
}

func (factory *multipartBodyFactory) SourceError(includeAbandoned bool) error {
	factory.mu.Lock()
	bodies := append([]*multipartBody(nil), factory.bodies...)
	factory.mu.Unlock()
	var result error
	for _, body := range bodies {
		if err := body.sourceError(includeAbandoned); err != nil {
			result = errors.Join(result, err)
		}
	}
	return result
}

func multipartBodyCloseErrors(bodies []*multipartBody) error {
	var result error
	for _, body := range bodies {
		if err := body.sourceError(false); err != nil {
			result = errors.Join(result, err)
		}
	}
	return result
}

type multipartBody struct {
	ctx        context.Context
	reader     io.Reader
	source     *multipartAudioSource
	stopCancel func() bool

	mu        sync.Mutex
	closed    bool
	abandoned bool
}

func (body *multipartBody) Read(buffer []byte) (int, error) {
	if err := body.ctx.Err(); err != nil {
		return 0, err
	}
	body.mu.Lock()
	closed := body.closed
	body.mu.Unlock()
	if closed {
		return 0, io.ErrClosedPipe
	}
	return body.reader.Read(buffer)
}

func (body *multipartBody) Close() error {
	body.mu.Lock()
	if body.closed {
		body.mu.Unlock()
		return body.source.Close()
	}
	body.closed = true
	body.abandoned = !body.source.reachedEOF()
	stopCancel := body.stopCancel
	body.mu.Unlock()
	if stopCancel != nil {
		stopCancel()
	}
	return body.source.Close()
}

func (body *multipartBody) sourceError(includeAbandoned bool) error {
	body.mu.Lock()
	abandoned := body.abandoned
	body.mu.Unlock()
	err := body.source.Err()
	if err != nil {
		return err
	}
	if abandoned && !includeAbandoned {
		return nil
	}
	return nil
}

type multipartAudioSource struct {
	audio io.ReadCloser

	mu          sync.Mutex
	closing     bool
	eof         bool
	readErr     error
	closeErr    error
	closeDone   chan struct{}
	closeResult error
	closed      bool
}

func (source *multipartAudioSource) Read(buffer []byte) (int, error) {
	n, err := source.audio.Read(buffer)
	source.mu.Lock()
	switch {
	case errors.Is(err, io.EOF):
		source.eof = true
	case err != nil && !source.closing:
		source.readErr = errors.Join(source.readErr, err)
	}
	source.mu.Unlock()
	return n, err
}

func (source *multipartAudioSource) Close() error {
	source.mu.Lock()
	if source.closed {
		done := source.closeDone
		source.mu.Unlock()
		<-done
		source.mu.Lock()
		err := source.closeResult
		source.mu.Unlock()
		return err
	}
	source.closed = true
	source.closing = true
	reachedEOF := source.eof
	source.mu.Unlock()

	err := source.audio.Close()
	source.mu.Lock()
	source.closeResult = err
	if err != nil && reachedEOF {
		source.closeErr = err
	}
	close(source.closeDone)
	source.mu.Unlock()
	return err
}

func (source *multipartAudioSource) reachedEOF() bool {
	source.mu.Lock()
	defer source.mu.Unlock()
	return source.eof
}

func (source *multipartAudioSource) Err() error {
	source.mu.Lock()
	defer source.mu.Unlock()
	return errors.Join(source.readErr, source.closeErr)
}