In a previous post I discussed asynchronous repositories. A closely related and complimentary design pattern is the Unit of Work pattern. In this post, I'll summarize the design pattern and cover a few non-conventional, but useful extensions.

Overview

The Unit of Work is a common design pattern used to manage the state changes to a set of objects. A unit of work abstracts all of the persistence operations and logic from other aspects of an application. Applying the pattern not only simplifies code that possess persistence needs, but it also makes changing or otherwise swapping out persistence strategies and methods easy.

A basic unit of work has the following characteristics:

  • Register New - registers an object for insertion.
  • Register Updated - registers an object for modification.
  • Register Removed - registers an object for deletion.
  • Commit - commits all pending work.
  • Rollback - discards all pending work.

Extensions

The basic design pattern supports most scenarios, but there are a few additional use cases that are typically not addressed. For stateful applications, it is usually desirable to support cancellation or simple undo operations by using deferred persistence. While this capability is covered via a rollback, there is not a way to interrogate whether a unit of work has pending changes.

Imagine your application has the following requirements:

  • As a user, I should only be able to save when there are uncommitted changes.
  • As a user, I should be prompted when I cancel an operation with uncommitted changes.

To satisfy these requirements, we only need to make a couple of additions:

  • Unregister - unregisters pending work for an object.
  • Has Pending Changes - indicates whether the unit of work contains uncommitted items.
  • Property Changed - raises an event when a property has changed.

Generic Interface

After reconsidering what is likely the majority of all plausible usage scenarios, we now have enough information to create a general purpose interface.

public interface IUnitOfWork<T> : INotifyPropertyChanged where T : class
{
    bool HasPendingChanges
    {
        get;
    }
    void RegisterNew( T item );
    void RegisterChanged( T item );
    void RegisterRemoved( T item );
    void Unregister( T item );
    void Rollback();
    Task CommitAsync( CancellationToken cancellationToken );
}

Base Implementation

It would be easy to stop at the generic interface definition, but we can do better. It is pretty straightforward to create a base implementation that handles just about everything except the commit operation.

public abstract class UnitOfWork<T> : IUnitOfWork<T> where T : class
{
    private readonly IEqualityComparer<T> comparer;
    private readonly HashSet<T> inserted;
    private readonly HashSet<T> updated;
    private readonly HashSet<T> deleted;

    protected UnitOfWork()
    protected UnitOfWork( IEqualityComparer<T> comparer )

    protected IEqualityComparer<T> Comparer { get; }
    protected virtual ICollection<T> InsertedItems { get; }
    protected virtual ICollection<T> UpdatedItems { get; }
    protected virtual ICollection<T> DeletedItems { get; }
    public virtual bool HasPendingChanges { get; }

    protected virtual void OnPropertyChanged( PropertyChangedEventArgs e )
    protected virtual void AcceptChanges()
    protected abstract bool IsNew( T item )
    public virtual void RegisterNew( T item )
    public virtual void RegisterChanged( T item )
    public virtual void RegisterRemoved( T item )
    public virtual void Unregister( T item )
    public virtual void Rollback()
    public abstract Task CommitAsync( CancellationToken cancellationToken );

    public event PropertyChangedEventHandler PropertyChanged;
}

Obviously by now, you've noticed that we've added a few protected members to support the implementation. We use HashSet<T> to track all inserts, updates, and deletes. By using HashSet<T>, we can easily ensure we don't track an entity more than once. We can also now apply some basic logic such as inserts should never enqueue for updates and deletes against uncommitted inserts should be negated. In addition, we add the ability to accept (e.g. clear) all pending work after the commit operation has completed successfully.

Supporting a Unit of Work Service Locator

Once we have all the previous pieces in place, we could again stop, but there are multiple ways in which a unit of work could be used in an application that we should consider:

  • Imperatively instantiated in code
  • Composed or inserted via dependency injection
  • Centrally retrieved via a special service locator facade

The decision as to which approach to use is at a developer's discretion. In general, when composition or dependency injection is used, the implementation is handed by another library and some mediating object (ex: a controller) will own the logic as to when or if entities are added to the unit of work. When a service locator is used, most or all of the logic can be baked directly into an object to enable self-tracking. In the rest of this section, we'll explore a UnitOfWork singleton that plays the role of a service locator.

public static class UnitOfWork
{
    public static IUnitOfWorkFactoryProvider Provider
    {
        get;
        set;
    }
    public static IUnitOfWork<TItem> Create<TItem>() where TItem : class
    public static IUnitOfWork<TItem> GetCurrent<TItem>() where TItem : class
    public static void SetCurrent<TItem>( IUnitOfWork<TItem> unitOfWork ) where TItem : class
    public static IUnitOfWork<TItem> NewCurrent<TItem>() where TItem : class
}

Populating the Service Locator

In order to locate a unit of work, the locator must be backed with code that can resolve it. We should also consider composite applications where there may be many units of work defined by different sources. The UnitOfWork singleton is configured by supplying an instance to the static Provider property.

Unit of Work Factory Provider

The IUnitOfWorkFactoryProvider interface can simply be thought of as a factory of factories. It provides a central mechanism for the service locator to resolve a unit of work via all known factories. In composite applications, implementers will likely want to use dependency injection. For ease of use, a default implementation is provided whose constructor accepts Func<IEnumerable<IUnitOfWorkFactory>>.

public interface IUnitOfWorkFactoryProvider
{
    IEnumerable<IUnitOfWorkFactory> Factories
    {
        get;
    }
}

Unit of Work Factory

The IUnitOfWorkFactory interface is used to register, create, and resolve units of work. Implementers have the option to map as many units of work to a factory as they like. In most scenarios, only one factory is required per application or composite component (ex: plug-in). A default implementation is provided that only requires the factory to register a function to create or resolve a unit of work for a given type. The Specification pattern is used to match or select the appropriate factory, but the exploration of that pattern is reserved for another time.

public interface IUnitOfWorkFactory
{
    ISpecification<Type> Specification
    {
        get;
    }
    IUnitOfWork<TItem> Create<TItem>() where TItem : class;
    IUnitOfWork<TItem> GetCurrent<TItem>() where TItem : class;
    void SetCurrent<TItem>( IUnitOfWork<TItem> unitOfWork ) where TItem : class;
}

Minimizing Test Setup

While all of factory interfaces make it flexible to support a configurable UnitOfWork singleton, it is somewhat painful to set up test cases. If the required unit of work is not resolved, an exception will be thrown; however, if the test doesn't involve a unit of work, why should we have to set one up?

To solve this problem, the service locator will internally create a compatible uncommitable unit of work instance whenever a unit of work cannot be resolved. This behavior allows self-tracking objects to be used without having to explicitly set up a mock or stub unit of work. You might be thinking that this behavior hides composition or dependency resolution failures and that is true. However, any attempt to commit against these instances will throw an InvalidOperationException, indicating that the unit of work is uncommitable. This approach is the most sensible method of avoiding unnecessary setups, while not completely hiding resolution failures. Whenever a unit of work fails in this manner, a developer should realize that they have not set up their test correctly (ex: verifying commit behavior) or resolution is failing at run time.

Examples

The following outlines some scenarios as to how a unit of work might be used. For each example, we'll use the following model:

public class Person
{
    public int PersonId { get; set;}
    public string FirstName { get; set; }
    public string LastName { get; set; }
}

Implementing a Unit of Work with the Entity Framework

The following demonstrates a simple unit of work that is backed by the Entity Framework:

public class PersonUnitOfWork : UnitOfWork<Person>
{
    protected override bool IsNew( Person item )
    {
        // any unsaved item will have an unset id
        return item.PersonId == 0;
    }
    public override async Task CommitAsync( CancellationToken cancellationToken )
    {
        using ( var context = new MyDbContext() )
        {
            foreach ( var item in this.Inserted )
                context.People.Add( item );

            foreach ( var item in this.Updated )
                context.People.Attach( item );

            foreach ( var item in this.Deleted )
                context.People.Remove( item );

            await context.SaveChangesAsync( cancellationToken );
        }
        this.AcceptChanges();
    }
}

Using a Unit of Work to Drive User Interactions

The following example illustrates using a unit of work in a rudimentary Windows Presentation Foundation (WPF) window that contains buttons to add, remove, cancel, and apply (or save) changes to a collection of people. The recommended approach to working with presentation layers such as WPF is to use the Model-View-View Model (MVVM) design pattern. For the sake of brevity and demonstration purposes, this example will use simple, albeit difficult to test, event handlers. All of the persistence logic is contained within the unit of work and the unit of work can report whether it has any pending work to help inform a user when there are changes. The unit of work can also be used to verify that the user truly wants to discard uncommitted changes, if there are any.

public partial class MyWindow : Window
{
    private readonly IUnitOfWork<Person> unitOfWork;
    public MyWindow() : this( new PersonUnitOfWork() ) { }
    public MyWindow( IUnitOfWork<Person> unitOfWork )
    {
        this.InitializeComponent();
        this.ApplyButton.IsEnabled = false;
        this.People = new ObservableCollection<Person>();
        this.unitOfWork = unitOfWork;
        this.unitOfWork.PropertyChanged +=
            ( s, e ) => this.ApplyButton.IsEnabled = this.unitOfWork.HasPendingChanges;
    }
    public Person SelectedPerson { get; set; }
    public ObservableCollection<Person> People { get; private set; }
    private void AddButton_Click( object sender, RoutedEventArgs e )
    {
        var person = new Person();
        // TODO: custom logic
        this.People.Add( person );
        this.unitOfWork.RegisterNew( person );
    }
    private void RemoveButton_Click( object sender, RoutedEventArgs e )
    {
        var person = this.SelectedPerson;
        if ( person == null ) return;
        this.People.Remove( person );
        this.unitOfWork.RegisterRemoved( person );
    }
    private async void ApplyButton_Click( object sender, RoutedEventArgs e )
    {
        await this.unitOfWork.CommitAsync( CancellationToken.None );
    }
    private void CancelButton_Click( object sender, RoutedEventArgs e )
    {
        if ( this.unitOfWork.HasPendingChanges )
        {
            var message = "Discard unsaved changes?";
            var title = "Save";
            var buttons = MessageBoxButton.YesNo;
            var answer = MessageBox.Show( message, title, buttons );
            if ( answer == DialogResult.No ) return;
            this.unitOfWork.Rollback();
        }
        this.Close();
    }
}

Implementing a Self-Tracking Entity

There are many different ways and varying degrees of functionality that can be implemented for a self-tracking entity. The following is one of many possibilities that illustrates just enough to convey the idea. The first thing we need to do is create a factory.

public class MyUnitOfWorkFactory : UnitOfWorkFactory
{
    public MyUnitOfWorkFactory()
    {
        this.RegisterFactoryMethod( () => new PersonUnitOfWork() );
        // additional units of work could be defined here
    }
}

Then we need to wire up the service locator with a provider that contains the factory.

var factories = new IUnitOfWorkFactory[]{ new MyUnitOfWorkFactory() };
UnitOfWork.Provider = new UnitOfWorkFactoryProvider( () => factories );

Finally, we can refactor the entity to enable self-tracking.

public class Person
{
    private string firstName;
    private string lastName;

    public int PersonId
    {
        get;
        set;
    }
    public string FirstName
    {
        get
        {
            return this.firstName;
        }
        set
        {
            this.firstName = value;
            UnitOfWork.GetCurrent<Person>().RegisterChanged( this );
        }
    }
    public string LastName
    {
        get
        {
            return this.lastName;
        }
        set
        {
            this.lastName = value;
            UnitOfWork.GetCurrent<Person>().RegisterChanged( this );
        }
    }
    public static Person CreateNew()
    {
        var person = new Person();
        UnitOfWork.GetCurrent<Person>().RegisterNew( person );
        return person;
    }
    public void Delete()
    {
        UnitOfWork.GetCurrent<Person>().RegisterRemoved( this );
    }
    public Task SaveAsync()
    {
        return UnitOfWork.GetCurrent<Person>().CommitAsync( CancellationToken.None );
    }
}

Conclusion

In this article we examined the Unit of Work pattern, added a few useful extensions to it, and demonstrated some common uses cases as to how you can apply the pattern. There are many implementations for the Unit of Work pattern and the concepts outlined in this article are no more correct than any of the alternatives. Hopefully you finish this article with a better understanding of the pattern and its potential uses. Although I didn't explicitly discuss unit testing, my belief is that most readers will recognize the benefits and ease in which cross-cutting persistence requirements can be tested using a unit of work. I've attached all the code required to leverage the Unit of Work pattern as described in this article in order to accelerate your own development, should you choose to do so.