diff --git a/client/single.go b/client/single.go new file mode 100644 index 0000000..bf0624d --- /dev/null +++ b/client/single.go @@ -0,0 +1,52 @@ +// Copyright 2019 Drone.IO Inc. All rights reserved. +// Use of this source code is governed by the Polyform License +// that can be found in the LICENSE file. + +package client + +import ( + "context" + "runtime/debug" + "sync" + + "github.com/drone/drone-go/drone" +) + +var _ Client = (*SingleFlight)(nil) + +// SingleFlight wraps a Client and limits to a single in-flight +// request to pull items from the queue. +type SingleFlight struct { + Client + mu sync.Mutex +} + +// NewSingleFlight returns a Client that is limited to a single in-flight +// request to pull items from the queue. +func NewSingleFlight(endpoint, secret string, skipverify bool) *SingleFlight { + return &SingleFlight{Client: New(endpoint, secret, skipverify)} +} + +// Request requests the next available build stage for execution. +func (t *SingleFlight) Request(ctx context.Context, args *Filter) (*drone.Stage, error) { + // if the context is canceled there is no need to make + // the request and we can exit early. + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + } + // if is critical to unlock the mutex when the function + // exits. although a panic is unlikely it is critical that + // we recover from the panic to avoid deadlock. + defer func() { + if r := recover(); r != nil { + debug.PrintStack() + } + t.mu.Unlock() + }() + // lock the mutex to ensure only a single in-flight + // request to request a resource from the server queue. + t.mu.Lock() + return t.Client.Request(ctx, args) +} diff --git a/client/single_test.go b/client/single_test.go new file mode 100644 index 0000000..d8454a1 --- /dev/null +++ b/client/single_test.go @@ -0,0 +1,77 @@ +// Copyright 2019 Drone.IO Inc. All rights reserved. +// Use of this source code is governed by the Polyform License +// that can be found in the LICENSE file. + +package client + +import ( + "context" + "errors" + "testing" + + "github.com/drone/drone-go/drone" +) + +var noContext = context.Background() + +func TestSingleFlight(t *testing.T) { + mock := &mockRequestClient{ + out: &drone.Stage{}, + err: errors.New("some random error"), + } + client := NewSingleFlight("", "", false) + client.Client = mock + out, err := client.Request(noContext, nil) + if got, want := out, mock.out; got != want { + t.Errorf("Expect stage returned from request") + } + if got, want := err, mock.err; got != want { + t.Errorf("Expect error returned from request") + } +} + +func TestSingleFlightPanic(t *testing.T) { + mock := &mockRequestClientPanic{} + client := NewSingleFlight("", "", false) + client.Client = mock + + defer func() { + if recover() != nil { + t.Errorf("Expect Request to recover from panic") + } + client.mu.Lock() + client.mu.Unlock() + }() + + client.Request(noContext, nil) +} + +func TestSingleFlightCancel(t *testing.T) { + ctx, cancel := context.WithCancel(noContext) + cancel() + client := NewSingleFlight("", "", false) + client.Request(ctx, nil) +} + +// mock client that returns a static stage and error +// from the request method. +type mockRequestClient struct { + Client + + out *drone.Stage + err error +} + +func (m *mockRequestClient) Request(ctx context.Context, args *Filter) (*drone.Stage, error) { + return m.out, m.err +} + +// mock client that returns panics when the request +// method is invoked. +type mockRequestClientPanic struct { + Client +} + +func (m *mockRequestClientPanic) Request(ctx context.Context, args *Filter) (*drone.Stage, error) { + panic("method not implemented") +}