Make WhisperX uploads replayable
This commit is contained in:
263
internal/adapters/whisperx/multipart_body.go
Normal file
263
internal/adapters/whisperx/multipart_body.go
Normal file
@@ -0,0 +1,263 @@
|
||||
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)
|
||||
}
|
||||
Reference in New Issue
Block a user