/***********************************************************************
Copyright (c) Microsoft Corporation
All Rights Reserved.
***********************************************************************/
// This source code is an intended supplement to the Microsoft SQL
// Server online references and related electronic documentation.
//
// This sample is for instructional purposes only.
// Code contained herein is not intended to be used "as is" in real applications.
// 
// mthread.cpp  : 
//  
// Test & demonstrate the use of multiple streams from a single process.
//
// This is a sample program used to demonstrate the Virtual Device Interface
// feature of Microsoft SQL Server.
//
// The program will backup or restore the 'pubs' sample database.
//
// The program requires two command line parameters.  
// 1)
// One of:
//  b   perform a backup
//  r   perform a restore
//
// 2)
// The number of streams to use, 1-32.
//  

#define _WIN32_DCOM

#include <objbase.h>    // for 'CoInitialize()'
#include <stdio.h>      // for file operations
#include <ctype.h>      // for toupper ()
#include <process.h>    // for C library-safe _beginthreadex,_endthreadex
#include <windows.h>

#include "vdi.h"        // interface declaration
#include "vdierror.h"   // error constants
#include "vdiguid.h"    // define the GUIDs 

void LogError (
    LPSTR           location,    // must always be provided
    LPSTR           description, // NULL is acceptable
    DWORD           errCode);    // windows status code

int performTransfer (
    IClientVirtualDevice*   vd,
    int                     backup,
    int                     streamId);

HANDLE execSQL (int doBackup, int nStreams);

int
startSecondaries(
    IClientVirtualDeviceSet *vds,
    HANDLE  hSQLProcess,  // handle to process dealing with the SQL
    int     nStreams);    // number of i/o streams

unsigned __stdcall
runSecondary (void *parms);

// Using a GUID for the VDS Name is a good way to assure uniqueness.
//
WCHAR	wVdsName [50];

//
// main function
//
int main(int argc, char *argv[])
{
    HRESULT                     hr;
    IClientVirtualDeviceSet*    vds = NULL ; 
    VDConfig                    config;
    int                         badParm=TRUE;
    int                         doBackup;
    HANDLE                      hProcess;
    int                         termCode = -1;
    int                         nStreams=1;

    // Check the input parm
    //
    if (argc == 3)
    {
        sscanf (argv[2], "%d", &nStreams);
        switch (toupper(argv[1][0]))
        {
            case 'B':
                doBackup = TRUE;
                badParm = FALSE;
                break;
        
            case 'R':
                doBackup = FALSE;
                badParm = FALSE;
                break;
        }
    }

    if (badParm)
    {
        printf ("useage: mprocess {B|R} <nStreams>\n"
            "Demonstrate a multistream Backup or Restore using the Virtual Device Interface\n");
        exit (1);
    }

    // 1..32 streams.
    //
    if (nStreams < 1)
            nStreams = 1;
    else if (nStreams > 32)
            nStreams = 32;

    printf ("Performing a %s using %d virtual device(s).\n", 
            (doBackup) ? "BACKUP" : "RESTORE", nStreams);

    // Initialize COM Library
    // Note: _WIN32_DCOM must be defined during the compile.
    //
    hr = CoInitializeEx (NULL, COINIT_MULTITHREADED);

    if (!SUCCEEDED (hr))
    {
        printf ("Coinit fails: x%X\n", hr);
        exit (1);
    }


    // Get an interface to the device set.
    // Notice how we use a single IID for both the class and interface
    // identifiers.
    //
    hr = CoCreateInstance ( 
        IID_IClientVirtualDeviceSet,
        NULL, 
        CLSCTX_INPROC_SERVER,
        IID_IClientVirtualDeviceSet,
        (void**)&vds);

    if (!SUCCEEDED (hr))
    {
        // This failure might happen if the DLL was not registered.
        //
        printf ("Could not create component: x%X\n", hr);
        printf ("Check registration of SQLVDI.DLL and value of IID\n");
        goto exit;
    }

	// Create a GUID to use for a unique virtual device name
	//
	GUID	vdsId;
	CoCreateGuid (&vdsId);
	StringFromGUID2 (vdsId, wVdsName, 49);

    // Setup the VDI configuration we want to use.
    // This program doesn't use any fancy features, so the
    // only field to setup is the deviceCount.
    //
    // The server will treat the virtual device just like a pipe:
    // I/O will be strictly sequential with only the basic commands.
    //
    memset (&config, 0, sizeof(config));

    config.deviceCount = nStreams;

    // Create the virtual device set
    //
    hr = vds->Create (wVdsName, &config);
    if (!SUCCEEDED (hr))
    {
        printf ("VDS::Create fails: x%X", hr);
        goto exit;
    }

    // Send the SQL command, via isql in a subprocess.
    //
    printf("\nSending the SQL...\n");

    hProcess = execSQL (doBackup, nStreams);
    if (hProcess == NULL)
    {
        printf ("execSQL failed.\n");
        goto shutdown;
    }


    // Wait for the server to connect, completing the configuration.
    // Notice that we wait a maximum of 15 seconds.
    //
    printf("\nWaiting for SQL to complete configuration...\n");

    hr = vds->GetConfiguration (15000, &config);
    if (!SUCCEEDED (hr))
    {
        printf ("VDS::Getconfig fails: x%X\n", hr);
        goto shutdown;
    }

    // Handle the virtual devices in secondary processes.
    //
    printf ("\nSpawning secondary threads..\n");
    termCode = startSecondaries (vds, hProcess, nStreams);
    
shutdown:

    // Close the set
    //
    vds->Close ();

    // COM reference counting: Release the interface.
    //
    vds->Release () ;

exit:

    // Uninitialize COM Library
    //
    CoUninitialize () ;

    return termCode;
}

//
// Execute a basic backup/restore, by starting 'isql' in a subprocess.
//
// Returns:
//  NULL    : failed to start isql
//  else    : process handle
//
HANDLE execSQL (int doBackup, int nStreams)
{
    char                    cmd[5000];
    char                    extend[100];
    PROCESS_INFORMATION     pi;
    STARTUPINFO             si;
    int                     ix;

	sprintf (cmd, "osql -E -b -Q\"%s DATABASE PUBS %s VIRTUAL_DEVICE='%ls'",
		(doBackup) ? "BACKUP" : "RESTORE",
		(doBackup) ? "TO" : "FROM",
		wVdsName);
    
    for (ix=1; ix<nStreams; ix++)
    {
        sprintf (extend, ", VIRTUAL_DEVICE='%ls%d'", wVdsName, ix);
        strcat (cmd, extend);
    }

    strcat (cmd, "\"");

    printf ("Submitting SQL:\n%s\n\n", cmd);

    // use my process for startup info
    //
    GetStartupInfo (&si);

    if (!CreateProcess (NULL, cmd, NULL, NULL,
            TRUE,   // inherit handles (stdin/stdout)
            0,      // creation flags,
            NULL, NULL,
            &si,    // startup info
            &pi))   // out: process info
    {
        LogError ("startSecondary", "CreateProcess", GetLastError ());
        return NULL;
    }
    
    // Return the process handle
    //
    return (pi.hProcess);
}

//-----------------------------------------------------------
//
// A parmameter block to use when spawning secondaries.
//
struct THREAD_PARMS {
    IClientVirtualDeviceSet *vds;
    int                      streamId;
};


//-----------------------------------------------------------
// Invoke the secondary threads, and wait for all children
// to complete.
//
// Returns: 0 if no errors were detected.
//
//
int
startSecondaries(
    IClientVirtualDeviceSet *vds,
    HANDLE  hSQLProcess,  // handle to process dealing with the SQL
    int     nStreams)     // number of i/o streams
{
    THREAD_PARMS    parms[32];      // each thread needs its own parm block
    int             ix,nActive;
    HANDLE          children[33];  // 32 is maximum number of streams.
                                   // plus one for the isql process.
    DWORD           waitStatus, exitCode;
    unsigned        threadId;

    for (ix=0; ix<nStreams; ix++)
    {
        // All threads share the same virtual device set,
        // but must operate on different virtual devices.
        //
        parms[ix].vds = vds;
        parms[ix].streamId = ix;

        children[ix] = (HANDLE)_beginthreadex (
            NULL, 0, runSecondary, (void*)&parms[ix], 0, &threadId);

        if (children[ix] == NULL)
        {
            printf ("Failed to create thread. errno is %d\n", errno);
            goto errorExit;
        }

        printf ("\nStarted thread %d\n", threadId);
    }

    // Add the isql process into the array
    //
    children[nStreams] = hSQLProcess;
    nActive = nStreams+1;

    // Wait for all to finish.
    // Max wait is one minute for this tiny test.
    //
    printf ("All children are now running.\n"
        "Waiting for their completion...\n");
    
    waitStatus = WaitForMultipleObjects (nActive, children, 
            TRUE, 60000);

    if (waitStatus < WAIT_OBJECT_0 ||
            waitStatus >= WAIT_OBJECT_0+nActive)
    {
        LogError ("startSecondary", "WaitForMultiple", GetLastError ());
        printf("Unexpected wait code: %d\n", waitStatus);
        goto errorExit;
    }

    // All of the children have completed.
    // Get the completion code from 'isql' to check for sucess.
    //
    if (!GetExitCodeProcess (hSQLProcess, &exitCode))
    {
        LogError ("startSecondary", "GetExitCode", GetLastError ());
        goto errorExit;
    }

    if (exitCode != 0)
    {
        printf ("The SQL operation failed with code %d\n", exitCode);
        goto errorExit;
    }

    printf ("The SQL operation was sucessful.\n");

    // Be sure to close handles when finished with them.
    //
    // Notice that in our trivial error handling here we
    // don't bother closing them, since handles are
    // automatically closed as part of process termination.
    //
    for (ix = 0; ix < nActive; ix++)
    {
        CloseHandle (children[ix]);
    }

    return 0;

errorExit:
    // Handle all problems in a trivial fashion:
    //  SignalAbort() will cause all processes using the virtual device set
    //  to terminate processing.
    //
    vds->SignalAbort ();

    //  However, since the threads are using the virtual device set allocated
    //  by the main thread, we can't let the main thread close the set before
    //  the threads are finished with it. There are two options:
    //  1) exit the process immediately.
    //  2) wait for the threads to terminate.
    //
    //  Option 2 is "cleaner" but requires more code, so we just exit here:
    //
    ExitProcess ((unsigned)-1);

    return -1;  // ExitProcess doesn't return; this avoids a compiler error.
}

//------------------------------------------------------------------
// Perform secondary client processing from within a thread.
// Exits with 0 if no errors detected, else nonzero.
// The caller must setup a THREAD_PARMS parameter block when
// spawning the thread.
//
unsigned __stdcall
runSecondary (void *parms)
{
    HRESULT                         hr;
    WCHAR                           devName[100];
    IClientVirtualDevice*           vd;
    VDConfig                        config;
    int                             termCode;

    // Fetch the input parms
    //
    int                         streamId = ((THREAD_PARMS*)parms)->streamId;
    IClientVirtualDeviceSet *   vds      = ((THREAD_PARMS*)parms)->vds;

    // Build the name of the device assigned to this thread.
    //
    if (streamId == 0)
    {
        // The first device has the same name as the set.
        //
        wcscpy (devName, wVdsName);
    }
    else
    {
        // For this example, we've simply appended a number
        // for additional devices.  You are free to name them
        // as you wish.
        //
        swprintf (devName, L"%s%d", wVdsName, streamId);
    }

    // Open the device assigned to this thread.
    //
    hr = vds->OpenDevice (devName, &vd);
    if (!SUCCEEDED (hr))
    {
        printf ("OpenDevice fails on %ls: x%X", devName, hr);
        termCode = -1;
        goto errExit;
    }

    // Grab the config to figure out data direction
    //
    hr = vds->GetConfiguration (INFINITE, &config);
    if (!SUCCEEDED (hr))
    {
        printf ("VDS::Getconfig fails: x%X\n", hr);
        termCode = -1;
        goto errExit;
    }

    printf ("\nPerforming data transfer...\n");
        
    termCode = performTransfer (vd, 
        (config.features&VDF_WriteMedia), streamId);

errExit:

    // If errors were detected, force an abort.
    //
    if (termCode != 0)
    {
        vds->SignalAbort ();
    }

    return ((unsigned)termCode);

}



// This routine reads commands from the server until a 'Close' status is received.
// It simply synchronously reads or writes a file on the root of the current drive.
//
// Returns 0, if no errors are detected, else non-zero.
//
int performTransfer (
    IClientVirtualDevice*   vd,
    int                     backup,
    int                     streamId)
{
    FILE *          fh;
    char            fname[80];
    VDC_Command *   cmd;
    DWORD           completionCode;
    DWORD           bytesTransferred;
    HRESULT         hr;
    int             termCode = -1;

    sprintf (fname, "multi.%d.dmp", streamId);

    fh = fopen (fname, (backup)? "wb" : "rb");
    if (fh == NULL )
    {
        printf ("Failed to open: %s\n", fname);
        return -1;
    }

    while (SUCCEEDED (hr=vd->GetCommand (INFINITE, &cmd)))
    {
        bytesTransferred = 0;
        switch (cmd->commandCode)
        {
            case VDC_Read:
                bytesTransferred = fread (cmd->buffer, 1, cmd->size, fh);
                if (bytesTransferred == cmd->size)
                    completionCode = ERROR_SUCCESS;
                else
                    // assume failure is eof
                    completionCode = ERROR_HANDLE_EOF;

                break;

            case VDC_Write:
                bytesTransferred = fwrite (cmd->buffer, 1, cmd->size, fh);
                if (bytesTransferred == cmd->size )
                {
                    completionCode = ERROR_SUCCESS;
                }
                else
                    // assume failure is disk full
                    completionCode = ERROR_DISK_FULL;
                break;

            case VDC_Flush:
                fflush (fh);
                completionCode = ERROR_SUCCESS;
                break;
    
            case VDC_ClearError:
                completionCode = ERROR_SUCCESS;
                break;

            default:
                // If command is unknown...
                completionCode = ERROR_NOT_SUPPORTED;
        }


        hr = vd->CompleteCommand (cmd, completionCode, bytesTransferred, 0);
        if (!SUCCEEDED (hr))
        {
            printf ("Completion Failed: x%X\n", hr);
            break;
        }
    }

    if (hr != VD_E_CLOSE)
    {
        printf ("Unexpected termination: x%X\n", hr);
    }
    else
    {
        // As far as the data transfer is concerned, no
        // errors occurred.  The code which issues the SQL
        // must determine if the backup/restore was
        // really successful.
        //
        printf ("Successfully completed data transfer.\n");
        termCode = 0;
    }

    fclose (fh);

    return termCode;
}


//--------------------------------------------------------------------
// 
// A simple error logger.
//
void LogError (
    LPSTR           location,    // must always be provided
    LPSTR           description, // NULL is acceptable
    DWORD           errCode)     // windows status code
{
    LPVOID lpMsgBuf;

    printf (
        "Error at %s: %s StatusCode: %X\n",
        location, 
        (description==NULL)?"":description,
        errCode);

    // Attempt to explain the code
    //
    if (errCode != 0 && FormatMessage(
        FORMAT_MESSAGE_ALLOCATE_BUFFER | 
            FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,    
        NULL,
        errCode,
        MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), // Default language
        (LPTSTR) &lpMsgBuf,    0,    NULL ) )// Process any inserts in lpMsgBuf.
    {
        printf ("Explanation: %s\n", lpMsgBuf);
        LocalFree( lpMsgBuf );
    }
}