Webdriver take screen shot when case failed use TestNG

Is there a good way to capture screenshots when running tests in parallel on the method level?
In order to run tests in parallel, each individual test needs a unique driver instance. So, at any given time you have x number of driver instances running. When it comes time to capture a screenshot, how do you determine which driver to use?

If you create a base test class with access to the driver, then that driver will always be the correct driver

The following will achieve this;

All test classes must extend a simple base test class;

public asbtract baseTestCase() {

private WebDriver driver;

public WebDriver getDriver() {
return driver;
}

@BeforeMethod
public void createDriver() {
Webdriver driver=XXXXDriver();
}

@AfterMethod
public void tearDownDriver() {
if (driver != null)
{
try
{
driver.quit();
}
catch (WebDriverException e) {
System.out.println("***** CAUGHT EXCEPTION IN DRIVER TEARDOWN *****");
System.out.println(e);
}

}
}

In your listener, you need to access the base class;

public class ScreenshotListener extends TestListenerAdapter {

@Override
public void onTestFailure(ITestResult result)
{
Object currentClass = result.getInstance();
WebDriver webDriver = ((BaseTest) currentClass).getDriver();

if (webDriver != null)
{

File f = ((TakesScreenshot) webDriver).getScreenshotAs(OutputType.FILE);

//etc.
}
}